<commit_before>// Copyright 2016 rust-fuzz developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

extern crate toml;
extern crate clap;
extern crate term;
#[macro_use]
extern crate error_chain;

use clap::{App, Arg, SubCommand, ArgMatches, AppSettings};
use std::{env, fs, path, ffi, process};
use std::io::Write;
use std::io::Read;

use std::os::unix::ffi::OsStrExt;

#[macro_use]
mod templates;
mod utils;

error_chain! {
    foreign_links {
        Toml(toml::de::Error);
        Io(::std::io::Error);
    }
}

fn main() {
    let app = App::new("cargo-fuzz")
        .version(option_env!("CARGO_PKG_VERSION").unwrap_or("0.0.0"))
        .about(option_env!("CARGO_PKG_DESCRIPTION").unwrap_or(""))
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .setting(AppSettings::GlobalVersion)
        // cargo passes in the subcommand name to the invoked executable. Use a hidden, optional
        // positional argument to deal with it?
        .arg(Arg::with_name("dummy").possible_value("fuzz").required(false).hidden(true))
        .subcommand(SubCommand::with_name("init").about("Initialize the fuzz folder"))
        .subcommand(SubCommand::with_name("run").about("Run the fuzz target in fuzz/fuzzers")
            .setting(AppSettings::TrailingVarArg)
            .arg(Arg::with_name("TARGET").required(true)
                 .help("name of the fuzz target"))
            .arg(Arg::with_name("ARGS").multiple(true)
                 .help("additional libFuzzer arguments passed to the binary \
                       (including a path to corpus or artefact which must be last argument)"))
        )
        .subcommand(SubCommand::with_name("add").about("Add a new fuzz target")
                    .arg(Arg::with_name("TARGET").required(true)
                         .help("name of the fuzz target"))
        )
        .subcommand(SubCommand::with_name("list").about("List all fuzz targets"));
    let args = app.get_matches();

    ::std::process::exit(match args.subcommand() {
        ("init", _) => FuzzProject::init().map(|_| ()),
        ("add", matches) =>
            FuzzProject::new().and_then(|p| p.add_target(matches.expect("arguments present"))),
        ("list", _) => FuzzProject::new().and_then(|p| p.list_targets()),
        ("run", matches) =>
            FuzzProject::new().and_then(|p| p.exec_target(matches.expect("arguments present"))),
        (s, _) => panic!("unimplemented subcommand {}!", s),
    }.map(|_| 0).unwrap_or_else(|err| {
        utils::report_error(err);
        1
    }));
}

struct FuzzProject {
    /// Path to the root cargo project
    ///
    /// Not the project with fuzz targets, but the project being fuzzed
    root_project: path::PathBuf,
    targets: Vec<String>,
}

impl FuzzProject {
    fn new() -> Result<Self> {
        let mut project = FuzzProject {
            root_project: find_package()?,
            targets: Vec::new()
        };
        let manifest = project.manifest()?;
        if !is_fuzz_manifest(&manifest) {
            return Err(format!("manifest `{:?}` does not look a cargo-fuzz manifest. \
                                Add following lines to override:\n\
                                [package.metadata]\ncargo-fuzz = true",
                                project.manifest_path()).into());
        }
        project.targets = collect_targets(&manifest);
        Ok(project)
    }

    /// Create the fuzz project structure
    ///
    /// This will not clone libfuzzer-sys
    fn init() -> Result<Self> {
        let project = FuzzProject {
            root_project: find_package()?,
            targets: Vec::new(),
        };
        let fuzz_project = project.path();
        // TODO: check if the project is already initialized
        fs::create_dir(&fuzz_project)?;
        fs::create_dir(fuzz_project.join("fuzzers"))?;

        let mut cargo = fs::File::create(fuzz_project.join("Cargo.toml"))?;
        cargo.write_fmt(toml_template!(project.root_project_name()?))?;

        let mut ignore = fs::File::create(fuzz_project.join(".gitignore"))?;
        ignore.write_fmt(gitignore_template!())?;

        const TARGET: &'static str = "fuzzer_script_1";
        project.create_target_template(TARGET)
               .chain_err(|| format!("could not create template file for target {:?}", TARGET))?;
        Ok(project)
    }

    fn list_targets(&self) -> Result<()> {
        for bin in &self.targets {
            utils::print_message(bin, term::color::GREEN);
        }
        Ok(())
    }

    fn add_target<'a>(&self, args: &ArgMatches<'a>) -> Result<()> {
        let target: String = args.value_of_os("TARGET").expect("TARGET is required").to_os_string()
            .into_string().map_err(|_| "TARGET must be valid unicode")?;
        // Create corpus and artifact directories for the newly added target
        self.corpus_for(&target)?;
        self.artifacts_for(&target)?;
        self.create_target_template(&target)
            .chain_err(|| format!("could not add target {:?}", target))
    }

    /// Add a new fuzz target script with a given name
    fn create_target_template<'a>(&self, target: &str) -> Result<()> {
        let target_path = self.target_path(target);
        let mut script = fs::OpenOptions::new().write(true).create_new(true).open(&target_path)
            .chain_err(|| format!("could not create target script file at {:?}", target_path))?;
        script.write_fmt(target_template!(self.root_project_name()?.replace("-", "_")))?;

        let mut cargo = fs::OpenOptions::new().append(true)
            .open(self.manifest_path())?;
        Ok(cargo.write_fmt(toml_bin_template!(target))?)
    }

    /// Fuzz a given fuzz target
    fn exec_target<'a>(&self, args: &ArgMatches<'a>) -> Result<()> {
        let target: String = args.value_of_os("TARGET").expect("TARGET is required").to_os_string()
            .into_string().map_err(|_| "TARGET must be valid unicode")?;
        let exec_args = args.values_of_os("ARGS")
                            .map(|v| v.collect::<Vec<_>>());
        let target_triple = "x86_64-unknown-linux-gnu";

        let mut flags = env::var("RUSTFLAGS").unwrap_or("".into());
        if !flags.is_empty() {
            flags.push(' ');
        }
        flags.push_str("-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");
        let mut cmd = process::Command::new("cargo");

        let mut artefact_arg = ffi::OsString::from("-artifact_prefix=");
        artefact_arg.push(self.artifacts_for(&target)?);

        // Merge the asan options, so users can still provide their own options
        // to e.g. disable the leak sanitizer. Options are colon-separated.
        let mut asan_opts = env::var("ASAN_OPTIONS").unwrap_or("".into());
        if !asan_opts.is_empty() {
            asan_opts.push(':');
        }
        asan_opts.push_str("detect_odr_violation=0");

        cmd.env("RUSTFLAGS", flags)
           .arg("run")
           .arg("--manifest-path")
           .arg(self.manifest_path())
           .arg("--verbose")
           .arg("--bin")
           .arg(&target)
           .arg("--target")
           // won't pass rustflags to build scripts
           .arg(target_triple)
           .arg("--")
           .arg(artefact_arg)
           .env("ASAN_OPTIONS", &asan_opts);

        // Corpus or artefact file must be last argument passed. Check if the last argument being
        // passed through to libfuzzer looks anything like a not-argument (missing `-`). If so,
        // do not append our pregenerated corpus directory.
        if let Some((last, rest)) = exec_args.as_ref().and_then(|args| args.split_last()) {
            for arg in rest {
                cmd.arg(arg);
            }
            if last.as_bytes().starts_with(b"-") {
                cmd.arg(last);
                cmd.arg(self.corpus_for(&target)?);
            } else {
                cmd.arg(last);
            }
        } else {
            cmd.arg(self.corpus_for(&target)?);
        }

        exec_cmd(&mut cmd).chain_err(|| format!("could not execute command: {:?}", cmd))?;
        Ok(())
    }

    fn path(&self) -> path::PathBuf {
        self.root_project.join("fuzz")
    }

    fn manifest_path(&self) -> path::PathBuf {
        self.path().join("Cargo.toml")
    }

    fn corpus_for(&self, target: &str) -> Result<path::PathBuf> {
        let mut p = self.path();
        p.push("corpus");
        p.push(target);
        fs::create_dir_all(&p)
            .chain_err(|| format!("could not make a corpus directory at {:?}", p))?;
        Ok(p)
    }

    fn artifacts_for(&self, target: &str) -> Result<path::PathBuf> {
        let mut p = self.path();
        p.push("artifacts");
        p.push(target);
        p.push(""); // trailing slash, necessary for libfuzzer, because it does simple concat
        fs::create_dir_all(&p)
            .chain_err(|| format!("could not make a artifact directory at {:?}", p))?;
        Ok(p)
    }

    fn target_path(&self, target: &str) -> path::PathBuf {
        let mut root = self.path();
        root.push("fuzzers");
        root.push(target);
        root.set_extension("rs");
        root
    }

    fn manifest(&self) -> Result<toml::Value> {
        let filename = self.manifest_path();
        let mut file = fs::File::open(&filename).chain_err(||
            format!("could not read the manifest file: {:?}", filename)
        )?;
        let mut data = Vec::new();
        file.read_to_end(&mut data)?;
        toml::from_slice(&data).chain_err(||
            format!("could not decode the manifest file at {:?}", filename)
        )
    }

    fn root_project_name(&self) -> Result<String> {
        let filename = self.root_project.join("Cargo.toml");
        let mut file = fs::File::open(&filename)?;
        let mut data = Vec::new();
        file.read_to_end(&mut data)?;
        let value: toml::Value = toml::from_slice(&data)?;
        let name = value.as_table().and_then(|v| v.get("package"))
                                   .and_then(toml::Value::as_table)
                                   .and_then(|v| v.get("name"))
                                   .and_then(toml::Value::as_str);
        if let Some(name) = name {
            Ok(String::from(name))
        } else {
            Err(format!("{:?} (package.name) is malformed", filename).into())
        }
    }
}

fn collect_targets(value: &toml::Value) -> Vec<String> {
    let bins = value.as_table().and_then(|v| v.get("bin"))
                               .and_then(toml::Value::as_array);
    if let Some(bins) = bins {
        bins.iter().map(|bin|
            bin.as_table().and_then(|v| v.get("name")).and_then(toml::Value::as_str)
        ).filter_map(|name| name.map(|v| String::from(v))).collect()
    } else {
        Vec::new()
    }
}

fn is_fuzz_manifest(value: &toml::Value) -> bool {
    let is_fuzz = value.as_table().and_then(|v| v.get("package"))
                                  .and_then(toml::Value::as_table)
                                  .and_then(|v| v.get("metadata"))
                                  .and_then(toml::Value::as_table)
                                  .and_then(|v| v.get("cargo-fuzz"))
                                  .and_then(toml::Value::as_bool);
    is_fuzz == Some(true)
}

/// Returns the path for the first found non-fuzz Cargo package
fn find_package() -> Result<path::PathBuf> {
    let mut dir = env::current_dir()?;
    let mut data = Vec::new();
    loop {
        let manifest_path = dir.join("Cargo.toml");
        match fs::File::open(&manifest_path) {
            Err(_) => {},
            Ok(mut f) => {
                f.read_to_end(&mut data)?;
                let value: toml::Value = toml::from_slice(&data)
                    .chain_err(||
                        format!("could not decode the manifest file at {:?}", manifest_path)
                    )?;
                if !is_fuzz_manifest(&value) {
                    // Not a cargo-fuzz project => must be a proper cargo project :)
                    return Ok(dir);
                }
            }
        }
        if !dir.pop() { break; }
    }
    Err("could not find a cargo project".into())
}

#[cfg(unix)]
fn exec_cmd(cmd: &mut process::Command) -> Result<process::ExitStatus> {
    use std::os::unix::process::CommandExt;
    Err(cmd.exec().into())
}

#[cfg(not(unix))]
fn exec_cmd(cmd: &mut process::Command) -> Result<process::ExitStatus> {
    cmd.status().map_err(|e| e.into())
}
<commit_msg>A flag to disable default corpus<commit_after>// Copyright 2016 rust-fuzz developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

extern crate toml;
extern crate clap;
extern crate term;
#[macro_use]
extern crate error_chain;

use clap::{App, Arg, SubCommand, ArgMatches, AppSettings};
use std::{env, fs, path, ffi, process};
use std::io::Write;
use std::io::Read;

#[macro_use]
mod templates;
mod utils;

error_chain! {
    foreign_links {
        Toml(toml::de::Error);
        Io(::std::io::Error);
    }
}

fn main() {
    let app = App::new("cargo-fuzz")
        .version(option_env!("CARGO_PKG_VERSION").unwrap_or("0.0.0"))
        .about(option_env!("CARGO_PKG_DESCRIPTION").unwrap_or(""))
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .setting(AppSettings::GlobalVersion)
        // cargo passes in the subcommand name to the invoked executable. Use a hidden, optional
        // positional argument to deal with it?
        .arg(Arg::with_name("dummy").possible_value("fuzz").required(false).hidden(true))
        .subcommand(SubCommand::with_name("init").about("Initialize the fuzz folder"))
        .subcommand(SubCommand::with_name("run").about("Run the fuzz target in fuzz/fuzzers")
            .arg(Arg::with_name("NO_CORPUS").short("n").long("no-corpus")
                 .help("do not automatically pass a corpus directory \
                       (allows for custom corpus or artefact files to be passed via <ARGS>)"))
            .arg(Arg::with_name("TARGET").required(true)
                 .help("name of the fuzz target"))
            .arg(Arg::with_name("ARGS").multiple(true)
                 .help("additional libFuzzer arguments passed to the binary"))
        )
        .subcommand(SubCommand::with_name("add").about("Add a new fuzz target")
                    .arg(Arg::with_name("TARGET").required(true)
                         .help("name of the fuzz target"))
        )
        .subcommand(SubCommand::with_name("list").about("List all fuzz targets"));
    let args = app.get_matches();

    ::std::process::exit(match args.subcommand() {
        ("init", _) => FuzzProject::init().map(|_| ()),
        ("add", matches) =>
            FuzzProject::new().and_then(|p| p.add_target(matches.expect("arguments present"))),
        ("list", _) => FuzzProject::new().and_then(|p| p.list_targets()),
        ("run", matches) =>
            FuzzProject::new().and_then(|p| p.exec_target(matches.expect("arguments present"))),
        (s, _) => panic!("unimplemented subcommand {}!", s),
    }.map(|_| 0).unwrap_or_else(|err| {
        utils::report_error(err);
        1
    }));
}

struct FuzzProject {
    /// Path to the root cargo project
    ///
    /// Not the project with fuzz targets, but the project being fuzzed
    root_project: path::PathBuf,
    targets: Vec<String>,
}

impl FuzzProject {
    fn new() -> Result<Self> {
        let mut project = FuzzProject {
            root_project: find_package()?,
            targets: Vec::new()
        };
        let manifest = project.manifest()?;
        if !is_fuzz_manifest(&manifest) {
            return Err(format!("manifest `{:?}` does not look a cargo-fuzz manifest. \
                                Add following lines to override:\n\
                                [package.metadata]\ncargo-fuzz = true",
                                project.manifest_path()).into());
        }
        project.targets = collect_targets(&manifest);
        Ok(project)
    }

    /// Create the fuzz project structure
    ///
    /// This will not clone libfuzzer-sys
    fn init() -> Result<Self> {
        let project = FuzzProject {
            root_project: find_package()?,
            targets: Vec::new(),
        };
        let fuzz_project = project.path();
        // TODO: check if the project is already initialized
        fs::create_dir(&fuzz_project)?;
        fs::create_dir(fuzz_project.join("fuzzers"))?;

        let mut cargo = fs::File::create(fuzz_project.join("Cargo.toml"))?;
        cargo.write_fmt(toml_template!(project.root_project_name()?))?;

        let mut ignore = fs::File::create(fuzz_project.join(".gitignore"))?;
        ignore.write_fmt(gitignore_template!())?;

        const TARGET: &'static str = "fuzzer_script_1";
        project.create_target_template(TARGET)
               .chain_err(|| format!("could not create template file for target {:?}", TARGET))?;
        Ok(project)
    }

    fn list_targets(&self) -> Result<()> {
        for bin in &self.targets {
            utils::print_message(bin, term::color::GREEN);
        }
        Ok(())
    }

    fn add_target<'a>(&self, args: &ArgMatches<'a>) -> Result<()> {
        let target: String = args.value_of_os("TARGET").expect("TARGET is required").to_os_string()
            .into_string().map_err(|_| "TARGET must be valid unicode")?;
        // Create corpus and artifact directories for the newly added target
        self.corpus_for(&target)?;
        self.artifacts_for(&target)?;
        self.create_target_template(&target)
            .chain_err(|| format!("could not add target {:?}", target))
    }

    /// Add a new fuzz target script with a given name
    fn create_target_template<'a>(&self, target: &str) -> Result<()> {
        let target_path = self.target_path(target);
        let mut script = fs::OpenOptions::new().write(true).create_new(true).open(&target_path)
            .chain_err(|| format!("could not create target script file at {:?}", target_path))?;
        script.write_fmt(target_template!(self.root_project_name()?.replace("-", "_")))?;

        let mut cargo = fs::OpenOptions::new().append(true)
            .open(self.manifest_path())?;
        Ok(cargo.write_fmt(toml_bin_template!(target))?)
    }

    /// Fuzz a given fuzz target
    fn exec_target<'a>(&self, args: &ArgMatches<'a>) -> Result<()> {
        let target: String = args.value_of_os("TARGET").expect("TARGET is required").to_os_string()
            .into_string().map_err(|_| "TARGET must be valid unicode")?;
        let no_corpus = args.is_present("NO_CORPUS");
        let exec_args = args.values_of_os("ARGS")
                            .map(|v| v.collect::<Vec<_>>());
        let target_triple = "x86_64-unknown-linux-gnu";

        let mut flags = env::var("RUSTFLAGS").unwrap_or("".into());
        if !flags.is_empty() {
            flags.push(' ');
        }
        flags.push_str("-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");
        let mut cmd = process::Command::new("cargo");

        let mut artefact_arg = ffi::OsString::from("-artifact_prefix=");
        artefact_arg.push(self.artifacts_for(&target)?);

        // Merge the asan options, so users can still provide their own options
        // to e.g. disable the leak sanitizer. Options are colon-separated.
        let mut asan_opts = env::var("ASAN_OPTIONS").unwrap_or("".into());
        if !asan_opts.is_empty() {
            asan_opts.push(':');
        }
        asan_opts.push_str("detect_odr_violation=0");

        cmd.env("RUSTFLAGS", flags)
           .arg("run")
           .arg("--manifest-path")
           .arg(self.manifest_path())
           .arg("--verbose")
           .arg("--bin")
           .arg(&target)
           .arg("--target")
           // won't pass rustflags to build scripts
           .arg(target_triple)
           .arg("--")
           .arg(artefact_arg)
           .env("ASAN_OPTIONS", &asan_opts);

        if let Some(args) = exec_args {
            for arg in args {
                cmd.arg(arg);
            }
        }
        if !no_corpus {
            cmd.arg(self.corpus_for(&target)?);
        }

        exec_cmd(&mut cmd).chain_err(|| format!("could not execute command: {:?}", cmd))?;
        Ok(())
    }

    fn path(&self) -> path::PathBuf {
        self.root_project.join("fuzz")
    }

    fn manifest_path(&self) -> path::PathBuf {
        self.path().join("Cargo.toml")
    }

    fn corpus_for(&self, target: &str) -> Result<path::PathBuf> {
        let mut p = self.path();
        p.push("corpus");
        p.push(target);
        fs::create_dir_all(&p)
            .chain_err(|| format!("could not make a corpus directory at {:?}", p))?;
        Ok(p)
    }

    fn artifacts_for(&self, target: &str) -> Result<path::PathBuf> {
        let mut p = self.path();
        p.push("artifacts");
        p.push(target);
        p.push(""); // trailing slash, necessary for libfuzzer, because it does simple concat
        fs::create_dir_all(&p)
            .chain_err(|| format!("could not make a artifact directory at {:?}", p))?;
        Ok(p)
    }

    fn target_path(&self, target: &str) -> path::PathBuf {
        let mut root = self.path();
        root.push("fuzzers");
        root.push(target);
        root.set_extension("rs");
        root
    }

    fn manifest(&self) -> Result<toml::Value> {
        let filename = self.manifest_path();
        let mut file = fs::File::open(&filename).chain_err(||
            format!("could not read the manifest file: {:?}", filename)
        )?;
        let mut data = Vec::new();
        file.read_to_end(&mut data)?;
        toml::from_slice(&data).chain_err(||
            format!("could not decode the manifest file at {:?}", filename)
        )
    }

    fn root_project_name(&self) -> Result<String> {
        let filename = self.root_project.join("Cargo.toml");
        let mut file = fs::File::open(&filename)?;
        let mut data = Vec::new();
        file.read_to_end(&mut data)?;
        let value: toml::Value = toml::from_slice(&data)?;
        let name = value.as_table().and_then(|v| v.get("package"))
                                   .and_then(toml::Value::as_table)
                                   .and_then(|v| v.get("name"))
                                   .and_then(toml::Value::as_str);
        if let Some(name) = name {
            Ok(String::from(name))
        } else {
            Err(format!("{:?} (package.name) is malformed", filename).into())
        }
    }
}

fn collect_targets(value: &toml::Value) -> Vec<String> {
    let bins = value.as_table().and_then(|v| v.get("bin"))
                               .and_then(toml::Value::as_array);
    if let Some(bins) = bins {
        bins.iter().map(|bin|
            bin.as_table().and_then(|v| v.get("name")).and_then(toml::Value::as_str)
        ).filter_map(|name| name.map(|v| String::from(v))).collect()
    } else {
        Vec::new()
    }
}

fn is_fuzz_manifest(value: &toml::Value) -> bool {
    let is_fuzz = value.as_table().and_then(|v| v.get("package"))
                                  .and_then(toml::Value::as_table)
                                  .and_then(|v| v.get("metadata"))
                                  .and_then(toml::Value::as_table)
                                  .and_then(|v| v.get("cargo-fuzz"))
                                  .and_then(toml::Value::as_bool);
    is_fuzz == Some(true)
}

/// Returns the path for the first found non-fuzz Cargo package
fn find_package() -> Result<path::PathBuf> {
    let mut dir = env::current_dir()?;
    let mut data = Vec::new();
    loop {
        let manifest_path = dir.join("Cargo.toml");
        match fs::File::open(&manifest_path) {
            Err(_) => {},
            Ok(mut f) => {
                f.read_to_end(&mut data)?;
                let value: toml::Value = toml::from_slice(&data)
                    .chain_err(||
                        format!("could not decode the manifest file at {:?}", manifest_path)
                    )?;
                if !is_fuzz_manifest(&value) {
                    // Not a cargo-fuzz project => must be a proper cargo project :)
                    return Ok(dir);
                }
            }
        }
        if !dir.pop() { break; }
    }
    Err("could not find a cargo project".into())
}

#[cfg(unix)]
fn exec_cmd(cmd: &mut process::Command) -> Result<process::ExitStatus> {
    use std::os::unix::process::CommandExt;
    Err(cmd.exec().into())
}

#[cfg(not(unix))]
fn exec_cmd(cmd: &mut process::Command) -> Result<process::ExitStatus> {
    cmd.status().map_err(|e| e.into())
}
<|endoftext|>


<commit_before>#![allow(non_snake_case)]

extern crate rand;
use rand::Rng;
use rand::distributions::{IndependentSample, Range};
mod energy;
use energy::*;
use std::io::prelude::*;
extern crate argparse;
use argparse::{ArgumentParser, Store, StoreFalse, StoreTrue};
mod trajectory;
use trajectory::*;

const LJ_EPS : f64 = 1.0;
const LJ_SIG : f64 = 1.0;

const TRIES_INTENDED : f64 = 3.0;
const DISP_SCALE_FACTOR : f64 = 0.1;

// easy printing to stderr
macro_rules! println_stderr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

fn parse_cmd_args(NUM_STEPS: &mut usize, NUM_MINIM_STEPS: &mut usize,
                  NUM_PARTICLES: &mut usize, DENSITY: &mut f64, TEMPERATURE: &mut f64,
                  CUTOFF: &mut f64, MAX_DISP_START: &mut f64, SCALE: &mut bool, TAILCORR: &mut bool, SHIFT: &mut bool,
                  OUTPUT_PREFIX: &mut String, OUTPUT_INTERVAL: &mut i64, OUTPUT_MINIM: &mut bool) {
    let mut ap = ArgumentParser::new();
    ap.set_description("LJ MC simulation.");
    ap.refer(NUM_STEPS)
        .add_option(&["-n", "--nsteps"], Store,
                    "Simulation steps: Number of steps for averaging" );
    ap.refer(NUM_MINIM_STEPS)
        .add_option(&["-m", "--nminimsteps"], Store,
                    "Minimization steps: Number of steps before averaging starts");
    ap.refer(NUM_PARTICLES)
        .add_option(&["-p", "--nparticles"], Store,
                    "Total number of particles");
    ap.refer(DENSITY)
        .add_option(&["-d", "--density"], Store,
                    "Particle density");
    ap.refer(TEMPERATURE)
        .add_option(&["-t", "--temperature"], Store,
                    "Temperature");
    ap.refer(CUTOFF)
        .add_option(&["--cutoff"], Store,
                    "Lennard jones cutoff radius in length of epsilon");
    ap.refer(MAX_DISP_START)
        .add_option(&["--displacement"], Store,
                    "Displacement per trial move");
    ap.refer(SCALE)
        .add_option(&["--nodisplacementscale"], StoreFalse,
                    "Disable displacement scaling");
    ap.refer(OUTPUT_PREFIX)
        .add_option(&["-o", "--output"], Store,
                    "Output file prefix");
    ap.refer(OUTPUT_INTERVAL)
        .add_option(&["--osteps"], Store,
                    "Number of steps between writing to the trajectory file. -1 only writes last frame");
    ap.refer(OUTPUT_MINIM)
        .add_option(&["--writeminimization"], StoreTrue,
                    "Enables writing of minimization step to trajectory");
//    ap.refer(VACUUM_DIMENSION)
//        .add_option(&["--vacuum"], Store,
//                    "Dimension of vacuum space below and above the intial system.");
    ap.refer(TAILCORR)
        .add_option(&["--notailcorr"], StoreFalse,
                    "Disable tailcorrection");
    ap.refer(SHIFT)
        .add_option(&["--noshift"], StoreFalse,
                    "Disable lj shifting");
    ap.parse_args_or_exit();
}

fn main() {

    // define all the stuff
    let mut minim_steps  = 1000000;
    let mut sample_steps = 100000;

    let mut num_particles: usize = 512;
    let mut density = 0.7;
    let mut temperature = 0.9;

    let mut cutoff = 3.0;
    let mut displacement = 0.1;

    let mut  TAILCORR : bool = true;
    let mut SHIFT: bool = true;
    let mut SCALE: bool = true;

    let mut output_prefix = "montecarlo".to_string();
    let mut output_interval : i64 = 100;
    let mut output_minim : bool = false;
    parse_cmd_args(&mut sample_steps, &mut minim_steps, &mut num_particles,
                   &mut density, &mut temperature,
                   &mut cutoff, &mut displacement, &mut SCALE, &mut TAILCORR, &mut SHIFT,
                   &mut output_prefix, &mut output_interval, &mut output_minim);

    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("##################  LJ Monte Carlo Simulation  #################");
    println_stderr!("################################################################");
    println_stderr!("");

    // initialize stuff
    let beta = 1.0/temperature;
    let volume = (num_particles as f64)/ density;
    let length  = volume.cbrt();
    let (l_x, l_y, l_z) = (length, length, length);
    let cutoff_squared = cutoff * cutoff;
    let max_displacement = length / 2.0;

    let mut rng = rand::thread_rng();
    let particle_range = Range::new(0, num_particles-1);

    let mut rx : Vec<f64> = vec![];
    let mut ry : Vec<f64> = vec![];
    let mut rz : Vec<f64> = vec![];
    loop {
        rx.push(l_x * rng.gen::<f64>());
        ry.push(l_y * rng.gen::<f64>());
        rz.push(l_z * rng.gen::<f64>());
        if rx.len() == num_particles { break; }
    }

    let e_shift = if SHIFT { 4.0 * LJ_EPS * ( (LJ_SIG/cutoff).powi(12) - (LJ_SIG/cutoff).powi(6) ) } else { 0.0 };
    let e_corr = if TAILCORR { 8.0/3.0*std::f64::consts::PI*density*LJ_EPS*LJ_SIG.powi(3)*((1.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };
    let p_corr = if TAILCORR { 16.0/3.0*std::f64::consts::PI*density.powi(2)*LJ_EPS*LJ_SIG.powi(3)*((2.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };

    println_stderr!("Particles: {}, Density: {}, Temperature: {}", num_particles, density, temperature);
    println_stderr!("System volume: {:8.3}, Dimensions {:.3}/{:.3}/{:.3}", volume, l_x, l_y, l_z);
    println_stderr!("Minimization steps: {}, Sampling steps: {}", minim_steps, sample_steps);
    println_stderr!("LJ params eps: {}, sigma: {}, cutoff: {}", LJ_EPS, LJ_SIG, cutoff);
    println_stderr!("Tailcorr: {:8.3}, Shift: {:8.3}, Pressurecprr: {:8.3}", e_corr, e_shift, p_corr);

    let (mut energy, mut virial) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
    let mut energy_sum = 0.0;
    let mut virial_sum = 0.0;
    let mut step_counter = 0;
    let mut accept_counter = 0;

    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("#####################  Energy Minimization  ####################");
    println_stderr!("################################################################");
    println_stderr!("");

    // prepare and write first trajectory frame
    let mut trajectory : XYZTrajectory = XYZTrajectory::new(&format!("{}.xyz", output_prefix));
    if output_minim { trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true); }

    for step in 0..minim_steps+sample_steps {

        // select rnd particle
        let rnd_index = particle_range.ind_sample(&mut rng);

        // store old position
        let oldX = rx[rnd_index];
        let oldY = ry[rnd_index];
        let oldZ = rz[rnd_index];

        // old particle energy
        let (old_particle_energy, old_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);

        // rnd displacement and PBC
        rx[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        ry[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        rz[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        if rx[rnd_index] < 0.0 { rx[rnd_index] += l_x }
        if rx[rnd_index] > l_x { rx[rnd_index] -= l_x }
        if ry[rnd_index] < 0.0 { ry[rnd_index] += l_y }
        if ry[rnd_index] > l_y { ry[rnd_index] -= l_y }
        if rz[rnd_index] < 0.0 { rz[rnd_index] += l_z }
        if rz[rnd_index] > l_z { rz[rnd_index] -= l_z }

        // calculate energy difference
        let (new_particle_energy, new_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);
        let dE = new_particle_energy - old_particle_energy;

        //accept move
        if rng.gen::<f64>() < (-beta * dE).exp() {
            accept_counter += 1;
            if step % 1000 == 0 { // calculate total energy every 1000 steps to account for rounding errors
                let (e, v) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
                energy = e;
                virial = v;
            } else {
                energy += dE;
                virial += new_particle_virial - old_particle_virial;
            }
        } else { // or restore old position
            rx[rnd_index] = oldX;
            ry[rnd_index] = oldY;
            rz[rnd_index] = oldZ;
        }

        // update sums for averaging
        step_counter += 1;
        energy_sum += energy;
        virial_sum += virial;

        // print some output during minimization
        if step < minim_steps && step_counter % 5000 == 0 && step != 0 {
            let tries_per_step : f64 = step_counter as f64 /accept_counter as f64;
            let acceptance_rate = 1.0/tries_per_step * 100.0;
            println_stderr!("Minim {:<10} Energy: {:<12.3} Virial: {:<12.3} Acceptance: {:<4.1}% Displacement: {:.3}", step, energy, virial, acceptance_rate, displacement);

            if SCALE {
                let scale_factor = (TRIES_INTENDED/tries_per_step * DISP_SCALE_FACTOR).abs();
                if tries_per_step < TRIES_INTENDED - 0.2 && displacement < max_displacement {
                    displacement += displacement * scale_factor;
                } else if tries_per_step > TRIES_INTENDED + 0.2 && displacement > 0.0 {
                    displacement -= displacement * scale_factor;
                }
                step_counter = 0;
                accept_counter = 0;
            }
        }

        // reset sums for sampling
        if step == minim_steps-1 {
            println_stderr!("");
            println_stderr!("################################################################");
            println_stderr!("##########################  Sampling  ##########################");
            println_stderr!("################################################################");
            println_stderr!("");
            step_counter = 0;
            accept_counter = 0;
            energy_sum = 0.0;
            virial_sum = 0.0;
        }

        if step > minim_steps && step_counter % 5000 == 0 {
            println_stderr!("Step  {:<10}Energy: {:<12.3}Virial: {:<12.3}", step_counter, energy, virial);
        }

        // write trajectory maybe
        if step as i64 % output_interval == 0 {
            if step > minim_steps || output_minim {
                trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
            }
        }

    }

    let final_energy = energy_sum/step_counter as f64;
    let particle_energy = final_energy / num_particles as f64;
    let final_virial = virial_sum / 3.0 / step_counter as f64 / volume;
    let pressure = virial_sum / 3.0 / step_counter as f64 / volume + density * temperature + p_corr;
    let final_acceptance_rate = 1.0/((accept_counter as f64)/(step_counter as f64)) * 100.0;

    println_stderr!("Done sampling!");
    println!("");
    println!("################################################################");
    println!("##########################  Results  ###########################");
    println!("################################################################");
    println!("");
    println!(
"Minimization: {}
Steps: {}

# Lennard Jones Params
epsilon: {}
sigma: {}
cutoff: {}

# System
Particles: {}
Density: {}
Temperature: {}
Volume: {}
Box dimension: {:.3}/{:.3}/{:.3}
Max Displacement: {}

# Correction
Energy correction: {}
Shift: {}
P-Correction: {}

# Averages
Tries: {}
Accepted: {}
Acceptance: {:.2}%
Energy: {}
Energy per particle: {}
Virial: {}
Pressure: {}",
         minim_steps, sample_steps,
        LJ_EPS, LJ_SIG, cutoff,
        num_particles, density, temperature, volume, l_x, l_y, l_z, displacement,
        e_corr, e_shift, p_corr,
        step_counter, accept_counter, final_acceptance_rate, final_energy, particle_energy, final_virial, pressure);

    trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
}
<commit_msg>vacuum slab<commit_after>#![allow(non_snake_case)]

extern crate rand;
use rand::Rng;
use rand::distributions::{IndependentSample, Range};
mod energy;
use energy::*;
use std::io::prelude::*;
extern crate argparse;
use argparse::{ArgumentParser, Store, StoreFalse, StoreTrue};
mod trajectory;
use trajectory::*;

const LJ_EPS : f64 = 1.0;
const LJ_SIG : f64 = 1.0;

const TRIES_INTENDED : f64 = 3.0;
const DISP_SCALE_FACTOR : f64 = 0.1;

// easy printing to stderr
macro_rules! println_stderr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

fn parse_cmd_args(NUM_STEPS: &mut usize, NUM_MINIM_STEPS: &mut usize,
                  NUM_PARTICLES: &mut usize, DENSITY: &mut f64, TEMPERATURE: &mut f64,
                  CUTOFF: &mut f64, MAX_DISP_START: &mut f64, SCALE: &mut bool, TAILCORR: &mut bool, SHIFT: &mut bool,
                  OUTPUT_PREFIX: &mut String, OUTPUT_INTERVAL: &mut i64, OUTPUT_MINIM: &mut bool,
                  VACUUM_SLAB: &mut f64) {
    let mut ap = ArgumentParser::new();
    ap.set_description("LJ MC simulation.");
    ap.refer(NUM_STEPS)
        .add_option(&["-n", "--nsteps"], Store,
                    "Simulation steps: Number of steps for averaging" );
    ap.refer(NUM_MINIM_STEPS)
        .add_option(&["-m", "--nminimsteps"], Store,
                    "Minimization steps: Number of steps before averaging starts");
    ap.refer(NUM_PARTICLES)
        .add_option(&["-p", "--nparticles"], Store,
                    "Total number of particles");
    ap.refer(DENSITY)
        .add_option(&["-d", "--density"], Store,
                    "Particle density");
    ap.refer(TEMPERATURE)
        .add_option(&["-t", "--temperature"], Store,
                    "Temperature");
    ap.refer(CUTOFF)
        .add_option(&["--cutoff"], Store,
                    "Lennard jones cutoff radius in length of epsilon");
    ap.refer(MAX_DISP_START)
        .add_option(&["--displacement"], Store,
                    "Displacement per trial move");
    ap.refer(SCALE)
        .add_option(&["--nodisplacementscale"], StoreFalse,
                    "Disable displacement scaling");
    ap.refer(OUTPUT_PREFIX)
        .add_option(&["-o", "--output"], Store,
                    "Output file prefix");
    ap.refer(OUTPUT_INTERVAL)
        .add_option(&["--osteps"], Store,
                    "Number of steps between writing to the trajectory file. -1 only writes last frame");
    ap.refer(OUTPUT_MINIM)
        .add_option(&["--writeminimization"], StoreTrue,
                    "Enables writing of minimization step to trajectory");
    ap.refer(VACUUM_SLAB)
        .add_option(&["--vacuum"], Store,
                    "Dimension of vacuum space above the intial system relative to the rest of the system (0=no slab, 1=half filled system, 2=thrid filled system ...).");
    ap.refer(TAILCORR)
        .add_option(&["--notailcorr"], StoreFalse,
                    "Disable tailcorrection");
    ap.refer(SHIFT)
        .add_option(&["--noshift"], StoreFalse,
                    "Disable lj shifting");
    ap.parse_args_or_exit();
}

fn main() {

    // define all the stuff
    let mut minim_steps  = 1000000;
    let mut sample_steps = 100000;

    let mut num_particles: usize = 512;
    let mut density = 0.7;
    let mut temperature = 0.9;

    let mut cutoff = 3.0;
    let mut displacement = 0.1;

    let mut  TAILCORR : bool = true;
    let mut SHIFT: bool = true;
    let mut SCALE: bool = true;

    let mut vacuum_slab = 0.0;

    let mut output_prefix = "montecarlo".to_string();
    let mut output_interval : i64 = 100;
    let mut output_minim : bool = false;
    parse_cmd_args(&mut sample_steps, &mut minim_steps, &mut num_particles,
                   &mut density, &mut temperature,
                   &mut cutoff, &mut displacement, &mut SCALE, &mut TAILCORR, &mut SHIFT,
                   &mut output_prefix, &mut output_interval, &mut output_minim,
                   &mut vacuum_slab);

    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("##################  LJ Monte Carlo Simulation  #################");
    println_stderr!("################################################################");
    println_stderr!("");

    // initialize stuff
    let beta = 1.0/temperature;
    let mut volume = (num_particles as f64)/ density;
    let length  = volume.cbrt();
    let (l_x, l_y, mut l_z) = (length, length, length);
    if vacuum_slab > 0.0 { // increase space in z
        let scale = vacuum_slab + 1.0;
        l_z *= scale;
        volume *= scale;
        density /= scale;
    }

    let cutoff_squared = cutoff * cutoff;
    let max_displacement = length / 2.0;

    let mut rng = rand::thread_rng();
    let particle_range = Range::new(0, num_particles-1);

    let mut rx : Vec<f64> = vec![];
    let mut ry : Vec<f64> = vec![];
    let mut rz : Vec<f64> = vec![];
    loop {
        rx.push(l_x * rng.gen::<f64>());
        ry.push(l_y * rng.gen::<f64>());
        rz.push(l_z * rng.gen::<f64>());
        if rx.len() == num_particles { break; }
    }

    let e_shift = if SHIFT { 4.0 * LJ_EPS * ( (LJ_SIG/cutoff).powi(12) - (LJ_SIG/cutoff).powi(6) ) } else { 0.0 };
    let e_corr = if TAILCORR { 8.0/3.0*std::f64::consts::PI*density*LJ_EPS*LJ_SIG.powi(3)*((1.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };
    let p_corr = if TAILCORR { 16.0/3.0*std::f64::consts::PI*density.powi(2)*LJ_EPS*LJ_SIG.powi(3)*((2.0/3.0*(LJ_SIG/cutoff).powi(9)) - (LJ_SIG/cutoff).powi(3)) } else { 0.0 };

    println_stderr!("Particles: {}, Density: {}, Temperature: {}", num_particles, density, temperature);
    println_stderr!("System volume: {:8.3}, Dimensions {:.3}/{:.3}/{:.3}", volume, l_x, l_y, l_z);
    println_stderr!("Minimization steps: {}, Sampling steps: {}", minim_steps, sample_steps);
    println_stderr!("LJ params eps: {}, sigma: {}, cutoff: {}", LJ_EPS, LJ_SIG, cutoff);
    println_stderr!("Tailcorr: {:8.3}, Shift: {:8.3}, Pressurecprr: {:8.3}", e_corr, e_shift, p_corr);

    let (mut energy, mut virial) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
    let mut energy_sum = 0.0;
    let mut virial_sum = 0.0;
    let mut step_counter = 0;
    let mut accept_counter = 0;

    println_stderr!("");
    println_stderr!("################################################################");
    println_stderr!("#####################  Energy Minimization  ####################");
    println_stderr!("################################################################");
    println_stderr!("");

    // prepare and write first trajectory frame
    let mut trajectory : XYZTrajectory = XYZTrajectory::new(&format!("{}.xyz", output_prefix));
    if output_minim { trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true); }

    for step in 0..minim_steps+sample_steps {

        // select rnd particle
        let rnd_index = particle_range.ind_sample(&mut rng);

        // store old position
        let oldX = rx[rnd_index];
        let oldY = ry[rnd_index];
        let oldZ = rz[rnd_index];

        // old particle energy
        let (old_particle_energy, old_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);

        // rnd displacement and PBC
        rx[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        ry[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        rz[rnd_index] += ( rng.gen::<f64>() - 0.5 ) * displacement;
        if rx[rnd_index] < 0.0 { rx[rnd_index] += l_x }
        if rx[rnd_index] > l_x { rx[rnd_index] -= l_x }
        if ry[rnd_index] < 0.0 { ry[rnd_index] += l_y }
        if ry[rnd_index] > l_y { ry[rnd_index] -= l_y }
        if rz[rnd_index] < 0.0 { rz[rnd_index] += l_z }
        if rz[rnd_index] > l_z { rz[rnd_index] -= l_z }

        // calculate energy difference
        let (new_particle_energy, new_particle_virial) = get_particle_energy(&rx, &ry, &rz, rnd_index, num_particles, l_x, l_y, l_z, cutoff_squared, e_shift);
        let dE = new_particle_energy - old_particle_energy;

        //accept move
        if rng.gen::<f64>() < (-beta * dE).exp() {
            accept_counter += 1;
            if step % 1000 == 0 { // calculate total energy every 1000 steps to account for rounding errors
                let (e, v) = get_total_energy(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, cutoff_squared, e_corr, e_shift);
                energy = e;
                virial = v;
            } else {
                energy += dE;
                virial += new_particle_virial - old_particle_virial;
            }
        } else { // or restore old position
            rx[rnd_index] = oldX;
            ry[rnd_index] = oldY;
            rz[rnd_index] = oldZ;
        }

        // update sums for averaging
        step_counter += 1;
        energy_sum += energy;
        virial_sum += virial;

        // print some output during minimization
        if step < minim_steps && step_counter % 5000 == 0 && step != 0 {
            let tries_per_step : f64 = step_counter as f64 /accept_counter as f64;
            let acceptance_rate = 1.0/tries_per_step * 100.0;
            println_stderr!("Minim {:<10} Energy: {:<12.3} Virial: {:<12.3} Acceptance: {:<4.1}% Displacement: {:.3}", step, energy, virial, acceptance_rate, displacement);

            if SCALE {
                let scale_factor = (TRIES_INTENDED/tries_per_step * DISP_SCALE_FACTOR).abs();
                if tries_per_step < TRIES_INTENDED - 0.2 && displacement < max_displacement {
                    displacement += displacement * scale_factor;
                } else if tries_per_step > TRIES_INTENDED + 0.2 && displacement > 0.0 {
                    displacement -= displacement * scale_factor;
                }
                step_counter = 0;
                accept_counter = 0;
            }
        }

        // reset sums for sampling
        if step == minim_steps-1 {
            println_stderr!("");
            println_stderr!("################################################################");
            println_stderr!("##########################  Sampling  ##########################");
            println_stderr!("################################################################");
            println_stderr!("");
            step_counter = 0;
            accept_counter = 0;
            energy_sum = 0.0;
            virial_sum = 0.0;
        }

        if step > minim_steps && step_counter % 5000 == 0 {
            println_stderr!("Step  {:<10}Energy: {:<12.3}Virial: {:<12.3}", step_counter, energy, virial);
        }

        // write trajectory maybe
        if step as i64 % output_interval == 0 {
            if step > minim_steps || output_minim {
                trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
            }
        }

    }

    let final_energy = energy_sum/step_counter as f64;
    let particle_energy = final_energy / num_particles as f64;
    let final_virial = virial_sum / 3.0 / step_counter as f64 / volume;
    let pressure = virial_sum / 3.0 / step_counter as f64 / volume + density * temperature + p_corr;
    let final_acceptance_rate = 1.0/((accept_counter as f64)/(step_counter as f64)) * 100.0;

    println_stderr!("Done sampling!");
    println!("");
    println!("################################################################");
    println!("##########################  Results  ###########################");
    println!("################################################################");
    println!("");
    println!(
"Minimization: {}
Steps: {}

# Lennard Jones Params
epsilon: {}
sigma: {}
cutoff: {}

# System
Particles: {}
Density: {}
Temperature: {}
Volume: {}
Box dimension: {:.3}/{:.3}/{:.3}
Max Displacement: {}

# Correction
Energy correction: {}
Shift: {}
P-Correction: {}

# Averages
Tries: {}
Accepted: {}
Acceptance: {:.2}%
Energy: {}
Energy per particle: {}
Virial: {}
Pressure: {}",
         minim_steps, sample_steps,
        LJ_EPS, LJ_SIG, cutoff,
        num_particles, density, temperature, volume, l_x, l_y, l_z, displacement,
        e_corr, e_shift, p_corr,
        step_counter, accept_counter, final_acceptance_rate, final_energy, particle_energy, final_virial, pressure);

    trajectory.write(&rx, &ry, &rz, num_particles, l_x, l_y, l_z, temperature, LJ_EPS, LJ_SIG, cutoff, true);
}
<|endoftext|>


<commit_before>// See LICENSE file for copyright and license details.

#[macro_use]
extern crate glium;

extern crate time;
extern crate image;
extern crate cgmath;

#[cfg(target_os = "android")]
extern crate android_glue;

mod fs;

use std::path::{Path};
use std::f32::consts::{PI};
use std::thread;
use std::time::Duration;
use glium::{glutin, Texture2d, DisplayBuild, Surface};
use glium::index::PrimitiveType;
use glium::glutin::ElementState::{Pressed, Released};
use cgmath::{Matrix4, Matrix3, Vector3, Vector2, Rad};

const FPS: u64 = 60;

#[derive(Copy, Clone)]
struct Vertex {
    position: [f32; 3],
    tex_coords: [f32; 2],
}

implement_vertex!(Vertex, position, tex_coords);

fn load_texture<P: AsRef<Path>>(display: &glium::Display, path: P) -> Texture2d {
    let f = fs::load(path);
    let image = image::load(f, image::PNG).unwrap().to_rgba();
    let image_dimensions = image.dimensions();
    let image = glium::texture::RawImage2d::from_raw_rgba(
        image.into_raw(), image_dimensions);
    Texture2d::new(display, image).unwrap()
}

fn make_program(display: &glium::Display) -> glium::Program {
    let api = display.get_window().unwrap().get_api();
    let pre_src = fs::load_string(match api {
        glutin::Api::OpenGl => "pre_gl.glsl",
        glutin::Api::OpenGlEs => "pre_gles.glsl",
        _ => unimplemented!(),
    });
    let vs_src = pre_src.clone() + &fs::load_string("vs.glsl");
    let fs_src = pre_src + &fs::load_string("fs.glsl");
    glium::Program::from_source(display, &vs_src, &fs_src, None).unwrap()
}

fn view_matrix(angle_x: Rad<f32>, angle_y: Rad<f32>, aspect: f32) -> Matrix4<f32> {
    let perspective_mat = cgmath::perspective(Rad::new(PI / 4.0), aspect, 0.1, 100.0);
    let tr_mat = Matrix4::from_translation(Vector3{x: 0.0, y: 0.0, z: -10.0});
    let angle_x_m = Matrix4::from(Matrix3::from_angle_z(angle_x));
    let angle_y_m = Matrix4::from(Matrix3::from_angle_x(angle_y));
    perspective_mat * tr_mat * angle_y_m * angle_x_m
}

fn win_size(display: &glium::Display) -> (u32, u32) {
    let window = display.get_window().unwrap();
    window.get_inner_size().unwrap()
}

fn aspect(display: &glium::Display) -> f32 {
    let (x, y) = win_size(display);
    x as f32 / y as f32
}

fn create_display() -> glium::Display {
    let gl_version = glutin::GlRequest::GlThenGles {
        opengles_version: (2, 0),
        opengl_version: (2, 1),
    };
    glutin::WindowBuilder::new()
        .with_gl(gl_version)
        .with_depth_buffer(24)
        .with_title("Nergal".to_string())
        .build_glium()
        .unwrap()
}

struct Visualizer {
    display: glium::Display,
    program: glium::Program,
    is_running: bool,
    accumulator: u64,
    previous_clock: u64,
    vertex_buffer: glium::VertexBuffer<Vertex>,
    index_buffer: glium::IndexBuffer<u16>,
    texture: Texture2d,
    camera_angle_x: Rad<f32>,
    camera_angle_y: Rad<f32>,
    aspect: f32,
    mouse_pos: Vector2<i32>,
    is_lmb_pressed: bool,
}

impl Visualizer {
    fn new() -> Visualizer {
        let display = create_display();
        let program = make_program(&display);
        let aspect = aspect(&display);
        let vertex_buffer = {
            let vertices = [
                Vertex { position: [-0.5, -0.5, 0.0], tex_coords: [0.0, 1.0] },
                Vertex { position: [-0.5,  0.5, 0.0], tex_coords: [0.0, 0.0] },
                Vertex { position: [ 0.5, -0.5, 0.0], tex_coords: [1.0, 1.0] },
                Vertex { position: [ 0.5,  0.5, 0.0], tex_coords: [1.0, 0.0] },
            ];
            glium::VertexBuffer::new(&display, &vertices).unwrap()
        };
        let index_buffer = glium::IndexBuffer::new(
            &display, PrimitiveType::TrianglesList, &[0u16, 1, 2, 1, 2, 3]).unwrap();
        let texture = load_texture(&display, "test.png");
        let accumulator = 0;
        let previous_clock = time::precise_time_ns();
        Visualizer {
            display: display,
            program: program,
            is_running: true,
            accumulator: accumulator,
            previous_clock: previous_clock,
            vertex_buffer: vertex_buffer,
            index_buffer: index_buffer,
            texture: texture,
            camera_angle_x: Rad::new(0.0),
            camera_angle_y: Rad::new(0.0),
            aspect: aspect,
            mouse_pos: Vector2{x: 0, y: 0},
            is_lmb_pressed: false,
        }
    }

    fn is_running(&self) -> bool {
        self.is_running
    }

    fn handle_events(&mut self) {
        self.aspect = aspect(&self.display);
        let rotate_step = PI / 12.0;
        let events: Vec<_> = self.display.poll_events().collect();
        for event in events {
            match event {
                glutin::Event::Resized(x, y) => {
                    self.aspect = x as f32 / y as f32;
                },
                glutin::Event::Closed => {
                    self.is_running = false;
                },
                glutin::Event::KeyboardInput(Released, _, Some(key)) => {
                    match key {
                        glutin::VirtualKeyCode::Q
                            | glutin::VirtualKeyCode::Escape =>
                        {
                            self.is_running = false;
                        },
                        glutin::VirtualKeyCode::Right => self.camera_angle_x.s += rotate_step,
                        glutin::VirtualKeyCode::Left => self.camera_angle_x.s -= rotate_step,
                        glutin::VirtualKeyCode::Down => self.camera_angle_y.s += rotate_step,
                        glutin::VirtualKeyCode::Up => self.camera_angle_y.s -= rotate_step,
                        _ => {},
                    }
                },
                glutin::Event::MouseMoved(x, y) => {
                    self.handle_mouse_move(Vector2{x: x, y: y});
                },
                glutin::Event::MouseInput(Pressed, glutin::MouseButton::Left) => {
                    self.is_lmb_pressed = true;
                },
                glutin::Event::MouseInput(Released, glutin::MouseButton::Left) => {
                    self.is_lmb_pressed = false;
                },
                glutin::Event::Touch(glutin::Touch{location: (x, y), phase, ..}) => {
                    let pos = Vector2{x: x as i32, y: y as i32};
                    match phase {
                        glutin::TouchPhase::Moved => {
                            self.handle_mouse_move(pos);
                        },
                        glutin::TouchPhase::Started => {
                            self.is_lmb_pressed = true;
                            self.mouse_pos = pos;
                        },
                        glutin::TouchPhase::Ended => {
                            self.mouse_pos = pos;
                            self.is_lmb_pressed = false;
                        },
                        glutin::TouchPhase::Cancelled => unimplemented!(),
                    }
                },
                _ => (),
            }
        }
    }

    fn handle_mouse_move(&mut self, pos: Vector2<i32>) {
        let diff = self.mouse_pos - pos;
        self.mouse_pos = pos;
        if self.is_lmb_pressed {
            let (w, h) = win_size(&self.display);
            self.camera_angle_x.s += PI * (diff.x as f32 / w as f32);
            self.camera_angle_y.s += PI * (diff.y as f32 / h as f32);
        }
    }

    fn draw(&self) {
        let view_mat: [[f32; 4]; 4] = view_matrix(
            self.camera_angle_x, self.camera_angle_y, self.aspect).into();
        let uniforms = uniform! {
            view_mat: view_mat,
            texture: &self.texture,
        };
        let mut target = self.display.draw();
        target.clear_color(0.0, 0.0, 0.0, 0.0);
        target.draw(
            &self.vertex_buffer,
            &self.index_buffer,
            &self.program,
            &uniforms,
            &Default::default(),
        ).unwrap();
        target.finish().unwrap();
    }

    fn update_timer(&mut self) {
        let fixed_time_stamp = 1_000_000_000 / FPS;
        let now = time::precise_time_ns();
        self.accumulator += now - self.previous_clock;
        self.previous_clock = now;
        while self.accumulator >= fixed_time_stamp {
            self.accumulator -= fixed_time_stamp;
            // TODO: update realtime state here
        }
        let remainder_ms = (fixed_time_stamp - self.accumulator) / 1_000_000;
        thread::sleep(Duration::from_millis(remainder_ms));
    }
}

fn main() {
    std::env::set_var("RUST_BACKTRACE", "1");
    let mut visualizer = Visualizer::new();
    while visualizer.is_running() {
        visualizer.draw();
        visualizer.handle_events();
        visualizer.update_timer();
    }
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<commit_msg>main.rs: Added basic model matrix<commit_after>// See LICENSE file for copyright and license details.

#[macro_use]
extern crate glium;

extern crate time;
extern crate image;
extern crate cgmath;

#[cfg(target_os = "android")]
extern crate android_glue;

mod fs;

use std::path::{Path};
use std::f32::consts::{PI};
use std::thread;
use std::time::Duration;
use glium::{glutin, Texture2d, DisplayBuild, Surface};
use glium::index::PrimitiveType;
use glium::glutin::ElementState::{Pressed, Released};
use cgmath::{Matrix4, Matrix3, Vector3, Vector2, Rad};

const FPS: u64 = 60;

#[derive(Copy, Clone)]
struct Vertex {
    position: [f32; 3],
    tex_coords: [f32; 2],
}

implement_vertex!(Vertex, position, tex_coords);

fn load_texture<P: AsRef<Path>>(display: &glium::Display, path: P) -> Texture2d {
    let f = fs::load(path);
    let image = image::load(f, image::PNG).unwrap().to_rgba();
    let image_dimensions = image.dimensions();
    let image = glium::texture::RawImage2d::from_raw_rgba(
        image.into_raw(), image_dimensions);
    Texture2d::new(display, image).unwrap()
}

fn make_program(display: &glium::Display) -> glium::Program {
    let api = display.get_window().unwrap().get_api();
    let pre_src = fs::load_string(match api {
        glutin::Api::OpenGl => "pre_gl.glsl",
        glutin::Api::OpenGlEs => "pre_gles.glsl",
        _ => unimplemented!(),
    });
    let vs_src = pre_src.clone() + &fs::load_string("vs.glsl");
    let fs_src = pre_src + &fs::load_string("fs.glsl");
    glium::Program::from_source(display, &vs_src, &fs_src, None).unwrap()
}

fn view_matrix(angle_x: Rad<f32>, angle_y: Rad<f32>, aspect: f32) -> Matrix4<f32> {
    let perspective_mat = cgmath::perspective(Rad::new(PI / 4.0), aspect, 0.1, 100.0);
    let tr_mat = Matrix4::from_translation(Vector3{x: 0.0, y: 0.0, z: -10.0});
    let angle_x_m = Matrix4::from(Matrix3::from_angle_z(angle_x));
    let angle_y_m = Matrix4::from(Matrix3::from_angle_x(angle_y));
    perspective_mat * tr_mat * angle_y_m * angle_x_m
}

fn win_size(display: &glium::Display) -> (u32, u32) {
    let window = display.get_window().unwrap();
    window.get_inner_size().unwrap()
}

fn aspect(display: &glium::Display) -> f32 {
    let (x, y) = win_size(display);
    x as f32 / y as f32
}

fn create_display() -> glium::Display {
    let gl_version = glutin::GlRequest::GlThenGles {
        opengles_version: (2, 0),
        opengl_version: (2, 1),
    };
    glutin::WindowBuilder::new()
        .with_gl(gl_version)
        .with_depth_buffer(24)
        .with_title("Nergal".to_string())
        .build_glium()
        .unwrap()
}

struct Visualizer {
    display: glium::Display,
    program: glium::Program,
    is_running: bool,
    accumulator: u64,
    previous_clock: u64,
    vertex_buffer: glium::VertexBuffer<Vertex>,
    index_buffer: glium::IndexBuffer<u16>,
    texture: Texture2d,
    camera_angle_x: Rad<f32>,
    camera_angle_y: Rad<f32>,
    aspect: f32,
    mouse_pos: Vector2<i32>,
    is_lmb_pressed: bool,
}

impl Visualizer {
    fn new() -> Visualizer {
        let display = create_display();
        let program = make_program(&display);
        let aspect = aspect(&display);
        let vertex_buffer = {
            let vertices = [
                Vertex { position: [-0.5, -0.5, 0.0], tex_coords: [0.0, 1.0] },
                Vertex { position: [-0.5,  0.5, 0.0], tex_coords: [0.0, 0.0] },
                Vertex { position: [ 0.5, -0.5, 0.0], tex_coords: [1.0, 1.0] },
                Vertex { position: [ 0.5,  0.5, 0.0], tex_coords: [1.0, 0.0] },
            ];
            glium::VertexBuffer::new(&display, &vertices).unwrap()
        };
        let index_buffer = glium::IndexBuffer::new(
            &display, PrimitiveType::TrianglesList, &[0u16, 1, 2, 1, 2, 3]).unwrap();
        let texture = load_texture(&display, "test.png");
        let accumulator = 0;
        let previous_clock = time::precise_time_ns();
        Visualizer {
            display: display,
            program: program,
            is_running: true,
            accumulator: accumulator,
            previous_clock: previous_clock,
            vertex_buffer: vertex_buffer,
            index_buffer: index_buffer,
            texture: texture,
            camera_angle_x: Rad::new(0.0),
            camera_angle_y: Rad::new(0.0),
            aspect: aspect,
            mouse_pos: Vector2{x: 0, y: 0},
            is_lmb_pressed: false,
        }
    }

    fn is_running(&self) -> bool {
        self.is_running
    }

    fn handle_events(&mut self) {
        self.aspect = aspect(&self.display);
        let rotate_step = PI / 12.0;
        let events: Vec<_> = self.display.poll_events().collect();
        for event in events {
            match event {
                glutin::Event::Resized(x, y) => {
                    self.aspect = x as f32 / y as f32;
                },
                glutin::Event::Closed => {
                    self.is_running = false;
                },
                glutin::Event::KeyboardInput(Released, _, Some(key)) => {
                    match key {
                        glutin::VirtualKeyCode::Q
                            | glutin::VirtualKeyCode::Escape =>
                        {
                            self.is_running = false;
                        },
                        glutin::VirtualKeyCode::Right => self.camera_angle_x.s += rotate_step,
                        glutin::VirtualKeyCode::Left => self.camera_angle_x.s -= rotate_step,
                        glutin::VirtualKeyCode::Down => self.camera_angle_y.s += rotate_step,
                        glutin::VirtualKeyCode::Up => self.camera_angle_y.s -= rotate_step,
                        _ => {},
                    }
                },
                glutin::Event::MouseMoved(x, y) => {
                    self.handle_mouse_move(Vector2{x: x, y: y});
                },
                glutin::Event::MouseInput(Pressed, glutin::MouseButton::Left) => {
                    self.is_lmb_pressed = true;
                },
                glutin::Event::MouseInput(Released, glutin::MouseButton::Left) => {
                    self.is_lmb_pressed = false;
                },
                glutin::Event::Touch(glutin::Touch{location: (x, y), phase, ..}) => {
                    let pos = Vector2{x: x as i32, y: y as i32};
                    match phase {
                        glutin::TouchPhase::Moved => {
                            self.handle_mouse_move(pos);
                        },
                        glutin::TouchPhase::Started => {
                            self.is_lmb_pressed = true;
                            self.mouse_pos = pos;
                        },
                        glutin::TouchPhase::Ended => {
                            self.mouse_pos = pos;
                            self.is_lmb_pressed = false;
                        },
                        glutin::TouchPhase::Cancelled => unimplemented!(),
                    }
                },
                _ => (),
            }
        }
    }

    fn handle_mouse_move(&mut self, pos: Vector2<i32>) {
        let diff = self.mouse_pos - pos;
        self.mouse_pos = pos;
        if self.is_lmb_pressed {
            let (w, h) = win_size(&self.display);
            self.camera_angle_x.s += PI * (diff.x as f32 / w as f32);
            self.camera_angle_y.s += PI * (diff.y as f32 / h as f32);
        }
    }

    fn draw(&self) {
        let model_pos = Vector3{x: 1.0, y: 0.0, z: 0.0};
        let model_mat: [[f32; 4]; 4] = Matrix4::from_translation(model_pos).into();
        let view_mat: [[f32; 4]; 4] = view_matrix(
            self.camera_angle_x, self.camera_angle_y, self.aspect).into();
        let uniforms = uniform! {
            view_mat: view_mat,
            model_mat: model_mat,
            texture: &self.texture,
        };
        let mut target = self.display.draw();
        target.clear_color(0.0, 0.0, 0.0, 0.0);
        target.draw(
            &self.vertex_buffer,
            &self.index_buffer,
            &self.program,
            &uniforms,
            &Default::default(),
        ).unwrap();
        target.finish().unwrap();
    }

    fn update_timer(&mut self) {
        let fixed_time_stamp = 1_000_000_000 / FPS;
        let now = time::precise_time_ns();
        self.accumulator += now - self.previous_clock;
        self.previous_clock = now;
        while self.accumulator >= fixed_time_stamp {
            self.accumulator -= fixed_time_stamp;
            // TODO: update realtime state here
        }
        let remainder_ms = (fixed_time_stamp - self.accumulator) / 1_000_000;
        thread::sleep(Duration::from_millis(remainder_ms));
    }
}

fn main() {
    std::env::set_var("RUST_BACKTRACE", "1");
    let mut visualizer = Visualizer::new();
    while visualizer.is_running() {
        visualizer.draw();
        visualizer.handle_events();
        visualizer.update_timer();
    }
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<|endoftext|>


<commit_before>#![feature(slice_patterns)]

extern crate nix;
extern crate ansi_term;
extern crate time;

use std::ffi::CString;
use std::io::{self, Write};
use nix::{unistd, sys};
use ansi_term::Colour::{Green, Red};
mod builtins;
mod commands;

use commands::Command;

fn main() {
    prompt(Ok(()));

    loop {
        let line = read_line().expect("can't read from stdin");
        let commands = commands::parse(&line);

        prompt(run(&commands));
    }
}

fn prompt(status: nix::Result<()>) {
    let time = time::now();
    let cwd = unistd::getcwd().unwrap();
    let icon = match status {
        Ok(_) => Green.paint("v"),
        Err(e) => Red.paint(format!("{}\nx", e.to_string())),
    };

    print!("{icon} {cwd} \n{time} $ ",
           icon=icon,
           time=time.strftime("%H:%M").unwrap(),
           cwd=cwd.to_str().unwrap());

    io::stdout().flush().unwrap()
}

fn read_line() -> io::Result<String> {
    let mut buffer = String::new();
    try!(io::stdin().read_line(&mut buffer));
    Ok(buffer)
}

fn run(commands: &[Command]) -> nix::Result<()> {
    for command in commands { exec(&command)? }
    sys::wait::wait().map(|_| ())
}

fn exec(command: &Command) -> nix::Result<()> {
    let result = try!(unistd::fork());
    if result.is_parent() { return Ok(()); }

    if command.fd.0 != 0 {
        unistd::dup2(command.fd.0, 0)?;
        unistd::close(command.fd.0)?;
    }

    if command.fd.1 != 1 {
        unistd::dup2(command.fd.1, 1)?;
        unistd::close(command.fd.1)?;
    }

    try!(match command.args.get(0) {
        Some(&"cd") => builtins::cd(&command.args),
        Some(&"echo") => builtins::echo(&command.args),
        _ => {
            let cargs: Vec<CString> = command.args.iter()
                .map(|s| CString::new(s.to_string()).unwrap())
                .collect();

            unistd::execvp(&cargs[0], &cargs)?;
            unreachable!()
        }
    });
    std::process::exit(0);
}
<commit_msg>Remove unused feature<commit_after>extern crate nix;
extern crate ansi_term;
extern crate time;

use std::ffi::CString;
use std::io::{self, Write};
use nix::{unistd, sys};
use ansi_term::Colour::{Green, Red};
mod builtins;
mod commands;

use commands::Command;

fn main() {
    prompt(Ok(()));

    loop {
        let line = read_line().expect("can't read from stdin");
        let commands = commands::parse(&line);

        prompt(run(&commands));
    }
}

fn prompt(status: nix::Result<()>) {
    let time = time::now();
    let cwd = unistd::getcwd().unwrap();
    let icon = match status {
        Ok(_) => Green.paint("v"),
        Err(e) => Red.paint(format!("{}\nx", e.to_string())),
    };

    print!("{icon} {cwd} \n{time} $ ",
           icon=icon,
           time=time.strftime("%H:%M").unwrap(),
           cwd=cwd.to_str().unwrap());

    io::stdout().flush().unwrap()
}

fn read_line() -> io::Result<String> {
    let mut buffer = String::new();
    try!(io::stdin().read_line(&mut buffer));
    Ok(buffer)
}

fn run(commands: &[Command]) -> nix::Result<()> {
    for command in commands { exec(&command)? }
    sys::wait::wait().map(|_| ())
}

fn exec(command: &Command) -> nix::Result<()> {
    let result = try!(unistd::fork());
    if result.is_parent() { return Ok(()); }

    if command.fd.0 != 0 {
        unistd::dup2(command.fd.0, 0)?;
        unistd::close(command.fd.0)?;
    }

    if command.fd.1 != 1 {
        unistd::dup2(command.fd.1, 1)?;
        unistd::close(command.fd.1)?;
    }

    try!(match command.args.get(0) {
        Some(&"cd") => builtins::cd(&command.args),
        Some(&"echo") => builtins::echo(&command.args),
        _ => {
            let cargs: Vec<CString> = command.args.iter()
                .map(|s| CString::new(s.to_string()).unwrap())
                .collect();

            unistd::execvp(&cargs[0], &cargs)?;
            unreachable!()
        }
    });
    std::process::exit(0);
}
<|endoftext|>


<commit_before>extern crate yaml_rust;
extern crate regex;

use std::fs::File;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(66);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(33);
        dotline.truncate(0);
      }
      else {
        result.push(' ');
        dotline.push(' ');
      }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 { result.push_str(&dotline); }
  result
}

fn main() {
  let mut app = App { print_ascii: false, print_binary: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]")
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut file = File::open("config.yml").unwrap();
  let mut config_str = String::new();
  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ascii");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  let app = app; // Revert to immutable

  let mut patterns = Vec::with_capacity(binary_matches.len());
  for &(_, pattern) in binary_matches.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = res.unwrap();
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);
          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string());
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut start = 0;
                  for i in 0..buf.len() {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        printables.push(&buf[start..i]);
                        found = false;
                      }
                    }
                  }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ascii", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", binary_matches[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string());
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<commit_msg>Give io timeouts more clear output<commit_after>extern crate yaml_rust;
extern crate regex;

use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::net::TcpListener;
use std::time::Duration;
use std::sync::Arc;
use std::thread;
use std::process::exit;
use yaml_rust::YamlLoader;
use regex::bytes::RegexSetBuilder;

#[derive(Copy, Clone)]
struct App {
  print_ascii: bool,
  print_binary: bool
}

fn to_hex(bytes: &[u8]) -> String {
  let mut count = 0;
  let mut result = String::with_capacity(67);
  let mut dotline = String::with_capacity(18);

  for byte in bytes.iter() {
    if count != 0 && count%8 == 0 {
      result.push(' ');
      if count%16 == 0 {
        dotline.push('\n');
        result.push_str(&dotline);
        result.reserve(67);
        dotline.truncate(0);
      }
      else { dotline.push(' '); }
    }
    result.push_str(&format!("{:02X} ", byte));
    if *byte > 31 && *byte < 127 { dotline.push(*byte as char) }
    else if *byte == 0 { dotline.push('-'); }
    else { dotline.push('.'); }
    count += 1;
  }
  while count%16 != 0 {
    if count%8 == 0 { result.push(' '); }
    result.push_str("   ");
    count += 1;
  }
  if dotline.len() != 0 { result.push_str(&dotline); }
  result
}

fn main() {
  let mut app = App { print_ascii: false, print_binary: false };
  let binary_matches = [
    ("SSL3.0 Record Protocol", r"^\x16\x03\x00..\x01"),
    ("TLS1.0 Record Protocol", r"^\x16\x03\x01..\x01"),
    ("TLS1.1 Record Protocol", r"^\x16\x03\x02..\x01"),
    ("TLS1.2 Record Protocol", r"^\x16\x03\x03..\x01"),
    ("SSL3.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x00"),
    ("TLS1.0 CLIENT_HELLO", r"^\x16....\x01...\x03\x01"),
    ("TLS1.1 CLIENT_HELLO", r"^\x16....\x01...\x03\x02"),
    ("TLS1.2 CLIENT_HELLO", r"^\x16....\x01...\x03\x03"),
    ("SMB1 COMMAND NEGOTIATE", r"^....\xffSMB\x72"),
    ("SMB1 NT_STATUS Success", r"^....\xffSMB.[\x00-\x0f]"),
    ("SMB1 NT_STATUS Information", r"^....\xffSMB.[\x40-\x4f]"),
    ("SMB1 NT_STATUS Warning", r"^....\xffSMB.[\x80-\x8f]"),
    ("SMB1 NT_STATUS Error", r"^....\xffSMB.[\xc0-\xcf]"),
    ("SMB2 COMMAND NEGOTIATE", r"^\x00...\xfeSMB........\x00\x00"),
    ("SMB2 NT_STATUS Success", r"^\x00...\xfeSMB....[\x00-\x0f]"),
    ("SMB2 NT_STATUS Information", r"^\x00...\xfeSMB....[\x40-\x4f]"),
    ("SMB2 NT_STATUS Warning", r"^\x00...\xfeSMB....[\x80-\x8f]"),
    ("SMB2 NT_STATUS Error", r"^\x00...\xfeSMB....[\xc0-\xcf]")
  ];
  let io_timeout = Duration::new(300, 0); // 5 minutes

  let mut file = File::open("config.yml").unwrap();
  let mut config_str = String::new();
  file.read_to_string(&mut config_str).unwrap();
  let docs = YamlLoader::load_from_str(&config_str).unwrap();
  let config = &docs[0];
  //println!("{:?}", config);
  if config["general"].is_badvalue() {
    println!("No 'general' section found in configuration file");
    exit(-1);
  }
  if config["ports"].is_badvalue() {
    println!("No 'ports' section found in configuration file");
    exit(-1);
  }
  if !config["general"]["print_ascii"].is_badvalue() {
    if config["general"]["print_ascii"].as_bool().unwrap() {
      app.print_ascii = true;
      println!("Printing ascii");
    }
  }
  if !config["general"]["print_binary"].is_badvalue() {
    if config["general"]["print_binary"].as_bool().unwrap() {
      app.print_binary = true;
      println!("Printing binary in hexadecimal");
    }
  }
  let app = app; // Revert to immutable

  let mut patterns = Vec::with_capacity(binary_matches.len());
  for &(_, pattern) in binary_matches.into_iter() {
    patterns.push(pattern);
  }
  let regexset = RegexSetBuilder::new(patterns)
    .unicode(false)
    .dot_matches_new_line(false)
    .build().unwrap();

  for port in config["ports"].as_vec().unwrap() {
    if !port["tcp"].is_badvalue() {
      let portno = port["tcp"].as_i64().unwrap();
      println!("TCP port {}", portno);
      let mut banner = Arc::new(String::new());
      if let Some(x) = port["banner"].as_str() {
        Arc::get_mut(& mut banner).unwrap().push_str(x);
        println!("  with banner: {}", *banner);
      }
      let regexset = regexset.clone();
      thread::spawn(move || {
        let server = TcpListener::bind(("0.0.0.0", portno as u16)).unwrap();
        for res in server.incoming() {
          let mut stream = res.unwrap();
          stream.set_read_timeout(Some(io_timeout)).expect("Failed to set read timeout on TcpStream");
          stream.set_write_timeout(Some(io_timeout)).expect("Failed to set write timeout on TcpStream");
          let addr = stream.peer_addr().unwrap();
          println!("CONNECT TCP {} from {}", portno, addr);
          let regexset = regexset.clone();
          let banner = banner.clone();
          thread::spawn(move || {
            if banner.len() > 0 {
              match stream.write((*banner).as_bytes()) {
                Ok(_) => {
                  for line in (*banner).lines() {
                    println!("> {}", line);
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("WRITE TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("WRITE ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  return;
                }
              }
            }
            let mut buf: [u8; 2048] = [0; 2048];
            loop {
              match stream.read(&mut buf) {
                Ok(c) => {
                  if c == 0 {
                    println!("CLOSE TCP {} from {}", portno, addr);
                    break;
                  }
                  let mut printables = Vec::new();
                  let mut found = false;
                  let mut start = 0;
                  for i in 0..buf.len() {
                    if (buf[i] > 31 && buf[i] < 127) || buf[i] == 10 || buf[i] == 13 {
                      if !found {
                        start = i;
                        found = true;
                      }
                    }
                    else {
                      if found {
                        printables.push(&buf[start..i]);
                        found = false;
                      }
                    }
                  }
                  if printables.len() == 1 && printables[0].len() == c {
                    if app.print_ascii {
                      let data = String::from_utf8_lossy(printables[0]);
                      for line in data.lines() {
                        println!("| {}", line);
                      }
                    }
                    else { println!("! Read {} bytes of printable ascii", c); }
                  }
                  else {
                    println!("! Read {} bytes of binary", c);
                    for id in regexset.matches(&buf[..c]).into_iter() {
                      println!("^ Matches pattern {}", binary_matches[id].0);
                    }
                    for printable in printables {
                      if printable.len() > 3 {
                        let data = String::from_utf8_lossy(printable);
                        for line in data.lines() {
                          println!("$ {}", line);
                        }
                      }
                    }
                    if app.print_binary {
                      let hex = to_hex(&buf[..c]);
                      for line in hex.lines() { println!(". {}", line); }
                    }
                  }
                }
                Err(e) => {
                  if e.kind() == io::ErrorKind::WouldBlock { println!("READ TIMEOUT TCP {} from {}", portno, addr); }
                  else { println!("READ ERROR TCP {} from {}: {}", portno, addr, e.to_string()); }
                  break;
                }
              }
              match stream.take_error() {
                Ok(opt) => {
                  if opt.is_some() {
                    println!("ERROR TCP {} from {}: {}", portno, addr, opt.unwrap().to_string());
                    break;
                  }
                }
                Err(_) => {
                  println!("This shouldn't happen...");
                  break;
                }
              }
            }
          });
        }
      });
    }
    else if !port["udp"].is_badvalue() {
      println!("UDP port {}", port["udp"].as_i64().unwrap());
    }
    else {
      println!("Invalid port specification in configuration file");
    }
  }

  loop { thread::sleep(Duration::new(60, 0)); }
}
<|endoftext|>