<commit_before>extern crate rustc_serialize;
extern crate docopt;
extern crate num;
extern crate flate2;
extern crate byteorder;
extern crate bincode;
extern crate rust_htslib as htslib;

mod tallyrun;
mod tallyread;
mod seqtable;
mod fasta;
mod filter;
mod counts;
mod bigwig;
mod scale;

use docopt::Docopt;
use std::process::exit;
use std::error::Error;
use std::path::Path;

/* Main usage/arguments */

const USAGE: &'static str = "
Cut-site frequencies

Usage:
  enzcut tallymer <fasta-file> <read-size> [--parts=<n>]
  enzcut seqtable <fasta-file> [options]
  enzcut dump <seqtbl-file> [<seqrange>]
  enzcut table <seqtbl-file> [<bam-file>...] [--qual=<q>] [--regions=<bedfile>] [--pdist=<min:max>] [--only-paired]
  enzcut scale <seqtbl-file> <bam-file>... [options]
  enzcut <fasta-file> <bam-file>... [options]
  enzcut (-h | --help)
  enzcut --version

Options:
  -h --help             Show this screen.
  --version             Show version.
  --cut-size=<n>        Cut-site size [default: 4].
  --tallymer=<file>     Unmappable positions file produced by tallymer (seq, pos).
  --plus-offset=<p>     Cut-site offset on plus strand [default: 2]. Eg, p=2 AA[A]A.
  --minus-offset=<m>    Cut-site offset on minus strand [default: 2]. Eg, m=2 A[A]AA.
  --read-size=<r>       Read length [default: 36].
  --parts=<n>           Split suffix tree generation into n parts [default: 4].
  --qual=<q>            Minimum read quality [default: 0].
  --regions=<bedfile>   Count only cut-sites inside the regions indicated in the BED file.
  --out=<outfile>       Output seqtable filename (defaults to fasta file basename with .tbl extension).
  --bed=<bedfile>       Output scaled BED filename (defaults to BAM file basename with '_scaled.bed' extension).
  --skip-bed            Skip creating the BED file output.
  --bw=<bigwigfile>     Output scaled BigWig filename (defaults to BAM file basename with .bw extension).
  --skip-bw             Skip creating the BigWig file output.
  --stranded            Output per strand counts when writting scaled values.
  --shift-counts        Shift minus strand counts.
  --no-scale            Skip actual scalling in 'scale' command.
  --pdist=<min:max>     Distance range for included paired reads.
  --only-paired         Only accept aligned reads that have a mapped pair.
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_fasta_file: String,
    arg_read_size: u16,
    arg_seqtbl_file: String,
    arg_seqrange: Option<String>,
    arg_bam_file: Option<Vec<String>>,
    flag_cut_size: u8,
    flag_tallymer: Option<String>,
    flag_plus_offset: u8,
    flag_minus_offset: u8,
    flag_version: bool,
    flag_read_size: u16,
    flag_parts: u8,
    flag_qual: u8,
    flag_regions: Option<String>,
    flag_out: Option<String>,
    flag_stranded: bool,
    flag_bed: Option<String>,
    flag_skip_bed: bool,
    flag_bw: Option<String>,
    flag_skip_bw: bool,
    flag_shift_counts: bool,
    flag_no_scale: bool,
    flag_pdist: Option<String>,
    flag_only_paired: bool,
    cmd_tallymer: bool,
    cmd_seqtable: bool,
    cmd_dump: bool,
    cmd_table: bool,
    cmd_scale: bool,
}

fn parse_range(range: &str) -> (i32, i32) {
    let parts: Vec<&str> = range.split(':').collect();

    if parts.len() != 2 {
        println!("Invalid distance range: {}", range);
        exit(1);
    }
    let min_val = parts[0].parse::<i32>();
    let max_val = parts[1].parse::<i32>();

    if min_val.is_err() || max_val.is_err() {
        println!("Invalid distance range: {}", range);
        exit(1);
    }
    (min_val.unwrap(),max_val.unwrap())
}

fn stem_filename(stem_src: &str, suffix: &str, out_arg: Option<String>) -> String {
    match out_arg {
        Some(out) => out,
        None => {
            let mut basename = Path::new(stem_src).file_stem().unwrap().to_os_string();
            basename.push(suffix);
            basename.into_string().unwrap()
        }
    }
}

fn main() {
    // Parse command line arguments
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("Cut-site frequencies, v{}.{}.{}",
            env!( "CARGO_PKG_VERSION_MAJOR" ),
            env!( "CARGO_PKG_VERSION_MINOR" ),
            env!( "CARGO_PKG_VERSION_PATCH" ) );
        return;
    }

    let dist_range = match args.flag_pdist {
        Some(range) => Some(parse_range(&range)),
        None => None
    };

    // Check for data output commands

    if args.cmd_dump {
        if let Some(seqrange) = args.arg_seqrange {
            seqtable::dump_seqtable_range(&args.arg_seqtbl_file, &seqrange);
        } else {
            seqtable::dump_seqtable(&args.arg_seqtbl_file);
        }

        return;
    }

    if args.cmd_table {
        let has_bam = args.arg_bam_file.is_some();
        let counts = counts::tabulate(&args.arg_seqtbl_file, args.arg_bam_file.as_ref(), args.flag_qual, args.flag_regions, dist_range, args.flag_only_paired);
        counts::print_counts(&counts, has_bam);
        return;
    }

    // Check for main sequence commands
    let mut run_tallymer = true;
    let mut run_seqtable = true;
    let mut run_scale = true;

    if args.cmd_tallymer {
        // just phase 1
        run_seqtable = false;
        run_scale = false;
    } else if args.cmd_seqtable {
        // phase 1 & 2
        run_scale = false;
    } else if args.cmd_scale {
        // just phase 3
        run_tallymer = false;
        run_seqtable = false;
    } else {
        // all three phases

        // catch cmd names being interpreted as fasta_file names
        if args.arg_fasta_file.eq("dump") || args.arg_fasta_file.eq("table") || args.arg_fasta_file.eq("tallymer") || args.arg_fasta_file.eq("seqtable") || args.arg_fasta_file.eq("scale") {
            println!("Invalid arguments to {} command.", args.arg_fasta_file);
            println!("{}", USAGE);
            exit(1);
        }
    }

    if args.flag_skip_bed && args.flag_skip_bw {
        println!("--skip-bed and --skip-bw cannot be used together");
        exit(1);
    }

    // Process main sequence phases

    // phase 1 - tallymer
    let tally_path = if run_tallymer {
        let path = tallyrun::tallymer_createfile(&args.arg_fasta_file, args.flag_read_size, args.flag_parts);
        println!("# tallymer produced/found {:}", path.to_str().unwrap());
        Some(path)
    } else { None };

    // phase 2 - seqtable
    let seqtable_file = if run_seqtable {
        let outfile = stem_filename(&args.arg_fasta_file, ".tbl", args.flag_out);

        let seq_params = seqtable::SeqTableParams {
            cut_length: args.flag_cut_size,
            plus_offset: args.flag_plus_offset,
            minus_offset: args.flag_minus_offset,
            read_length: args.flag_read_size };
        fasta::process_fasta(&args.arg_fasta_file, &tally_path.unwrap(), seq_params, &outfile);
        println!("# seqtable produced {}", &outfile);
        outfile
    } else {
        args.arg_seqtbl_file
    };

    // phase 3 - tabulate & scale
    if run_scale {
        let bamfile = args.arg_bam_file.as_ref().unwrap()[0].clone(); // use the first name for reference
        let counts = counts::tabulate(&seqtable_file, args.arg_bam_file.as_ref(), args.flag_qual, args.flag_regions, dist_range, args.flag_only_paired);
        let pileup = scale::scale(&seqtable_file, counts, args.arg_bam_file.as_ref().unwrap(), args.flag_qual, args.flag_shift_counts, args.flag_no_scale, &dist_range, args.flag_only_paired);

        if !args.flag_skip_bed {
            let outfile_bed = if args.flag_no_scale {
                    stem_filename(&bamfile, "_not_scaled.bed", args.flag_bed)
                } else {
                    stem_filename(&bamfile, "_scaled.bed", args.flag_bed)
                };

            match pileup.write_bed(&outfile_bed, args.flag_stranded) {
                Ok(_) => println!("# scale produced {}", &outfile_bed),
                Err(err) => println!("Error producing BED file: {}", err.description()),
            }
        } else {
            println!("# scale skipping BED output");
        }

        if !args.flag_skip_bw {
            let outfile_bw = stem_filename(&bamfile, ".bw", args.flag_bw);

            match pileup.write_bw(&outfile_bw, args.flag_stranded) {
                Ok((f1, f2)) => {
                    if f2.is_some() {
                        println!("# scale produced {}", f1);
                        println!("# scale produced {}", f2.unwrap());
                    } else {
                        println!("# scale produced {}", f1);
                    }
                },
                Err(err) => println!("Error producing BigWig file: {}", err.description()),
            }
        } else {
            println!("# scale skipping BigWig output");
        }
    }
}
<commit_msg>Change default seqtable output filename to include parameter values<commit_after>extern crate rustc_serialize;
extern crate docopt;
extern crate num;
extern crate flate2;
extern crate byteorder;
extern crate bincode;
extern crate rust_htslib as htslib;

mod tallyrun;
mod tallyread;
mod seqtable;
mod fasta;
mod filter;
mod counts;
mod bigwig;
mod scale;

use docopt::Docopt;
use std::process::exit;
use std::error::Error;
use std::path::Path;

/* Main usage/arguments */

const USAGE: &'static str = "
Cut-site frequencies

Usage:
  enzcut tallymer <fasta-file> <read-size> [--parts=<n>]
  enzcut seqtable <fasta-file> [options]
  enzcut dump <seqtbl-file> [<seqrange>]
  enzcut table <seqtbl-file> [<bam-file>...] [--qual=<q>] [--regions=<bedfile>] [--pdist=<min:max>] [--only-paired]
  enzcut scale <seqtbl-file> <bam-file>... [options]
  enzcut <fasta-file> <bam-file>... [options]
  enzcut (-h | --help)
  enzcut --version

Options:
  -h --help             Show this screen.
  --version             Show version.
  --cut-size=<n>        Cut-site size [default: 4].
  --tallymer=<file>     Unmappable positions file produced by tallymer (seq, pos).
  --plus-offset=<p>     Cut-site offset on plus strand [default: 2]. Eg, p=2 AA[A]A.
  --minus-offset=<m>    Cut-site offset on minus strand [default: 2]. Eg, m=2 A[A]AA.
  --read-size=<r>       Read length [default: 36].
  --parts=<n>           Split suffix tree generation into n parts [default: 4].
  --qual=<q>            Minimum read quality [default: 0].
  --regions=<bedfile>   Count only cut-sites inside the regions indicated in the BED file.
  --out=<outfile>       Output seqtable filename (defaults to fasta file basename with .tbl extension).
  --bed=<bedfile>       Output scaled BED filename (defaults to BAM file basename with '_scaled.bed' extension).
  --skip-bed            Skip creating the BED file output.
  --bw=<bigwigfile>     Output scaled BigWig filename (defaults to BAM file basename with .bw extension).
  --skip-bw             Skip creating the BigWig file output.
  --stranded            Output per strand counts when writting scaled values.
  --shift-counts        Shift minus strand counts.
  --no-scale            Skip actual scalling in 'scale' command.
  --pdist=<min:max>     Distance range for included paired reads.
  --only-paired         Only accept aligned reads that have a mapped pair.
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_fasta_file: String,
    arg_read_size: u16,
    arg_seqtbl_file: String,
    arg_seqrange: Option<String>,
    arg_bam_file: Option<Vec<String>>,
    flag_cut_size: u8,
    flag_tallymer: Option<String>,
    flag_plus_offset: u8,
    flag_minus_offset: u8,
    flag_version: bool,
    flag_read_size: u16,
    flag_parts: u8,
    flag_qual: u8,
    flag_regions: Option<String>,
    flag_out: Option<String>,
    flag_stranded: bool,
    flag_bed: Option<String>,
    flag_skip_bed: bool,
    flag_bw: Option<String>,
    flag_skip_bw: bool,
    flag_shift_counts: bool,
    flag_no_scale: bool,
    flag_pdist: Option<String>,
    flag_only_paired: bool,
    cmd_tallymer: bool,
    cmd_seqtable: bool,
    cmd_dump: bool,
    cmd_table: bool,
    cmd_scale: bool,
}

fn parse_range(range: &str) -> (i32, i32) {
    let parts: Vec<&str> = range.split(':').collect();

    if parts.len() != 2 {
        println!("Invalid distance range: {}", range);
        exit(1);
    }
    let min_val = parts[0].parse::<i32>();
    let max_val = parts[1].parse::<i32>();

    if min_val.is_err() || max_val.is_err() {
        println!("Invalid distance range: {}", range);
        exit(1);
    }
    (min_val.unwrap(),max_val.unwrap())
}

fn stem_filename(stem_src: &str, suffix: &str, out_arg: Option<String>) -> String {
    match out_arg {
        Some(out) => out,
        None => {
            let mut basename = Path::new(stem_src).file_stem().unwrap().to_os_string();
            basename.push(suffix);
            basename.into_string().unwrap()
        }
    }
}

fn main() {
    // Parse command line arguments
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("Cut-site frequencies, v{}.{}.{}",
            env!( "CARGO_PKG_VERSION_MAJOR" ),
            env!( "CARGO_PKG_VERSION_MINOR" ),
            env!( "CARGO_PKG_VERSION_PATCH" ) );
        return;
    }

    let dist_range = match args.flag_pdist {
        Some(range) => Some(parse_range(&range)),
        None => None
    };

    // Check for data output commands

    if args.cmd_dump {
        if let Some(seqrange) = args.arg_seqrange {
            seqtable::dump_seqtable_range(&args.arg_seqtbl_file, &seqrange);
        } else {
            seqtable::dump_seqtable(&args.arg_seqtbl_file);
        }

        return;
    }

    if args.cmd_table {
        let has_bam = args.arg_bam_file.is_some();
        let counts = counts::tabulate(&args.arg_seqtbl_file, args.arg_bam_file.as_ref(), args.flag_qual, args.flag_regions, dist_range, args.flag_only_paired);
        counts::print_counts(&counts, has_bam);
        return;
    }

    // Check for main sequence commands
    let mut run_tallymer = true;
    let mut run_seqtable = true;
    let mut run_scale = true;

    if args.cmd_tallymer {
        // just phase 1
        run_seqtable = false;
        run_scale = false;
    } else if args.cmd_seqtable {
        // phase 1 & 2
        run_scale = false;
    } else if args.cmd_scale {
        // just phase 3
        run_tallymer = false;
        run_seqtable = false;
    } else {
        // all three phases

        // catch cmd names being interpreted as fasta_file names
        if args.arg_fasta_file.eq("dump") || args.arg_fasta_file.eq("table") || args.arg_fasta_file.eq("tallymer") || args.arg_fasta_file.eq("seqtable") || args.arg_fasta_file.eq("scale") {
            println!("Invalid arguments to {} command.", args.arg_fasta_file);
            println!("{}", USAGE);
            exit(1);
        }
    }

    if args.flag_skip_bed && args.flag_skip_bw {
        println!("--skip-bed and --skip-bw cannot be used together");
        exit(1);
    }

    // Process main sequence phases

    // phase 1 - tallymer
    let tally_path = if run_tallymer {
        let path = tallyrun::tallymer_createfile(&args.arg_fasta_file, args.flag_read_size, args.flag_parts);
        println!("# tallymer produced/found {:}", path.to_str().unwrap());
        Some(path)
    } else { None };

    // phase 2 - seqtable
    let seqtable_file = if run_seqtable {
        let suffix = format!("_{}.{}.{}.{}.tbl", args.flag_read_size, args.flag_cut_size, args.flag_plus_offset, args.flag_minus_offset);
        let outfile = stem_filename(&args.arg_fasta_file, &suffix, args.flag_out);

        let seq_params = seqtable::SeqTableParams {
            cut_length: args.flag_cut_size,
            plus_offset: args.flag_plus_offset,
            minus_offset: args.flag_minus_offset,
            read_length: args.flag_read_size };
        fasta::process_fasta(&args.arg_fasta_file, &tally_path.unwrap(), seq_params, &outfile);
        println!("# seqtable produced {}", &outfile);
        outfile
    } else {
        args.arg_seqtbl_file
    };

    // phase 3 - tabulate & scale
    if run_scale {
        let bamfile = args.arg_bam_file.as_ref().unwrap()[0].clone(); // use the first name for reference
        let counts = counts::tabulate(&seqtable_file, args.arg_bam_file.as_ref(), args.flag_qual, args.flag_regions, dist_range, args.flag_only_paired);
        let pileup = scale::scale(&seqtable_file, counts, args.arg_bam_file.as_ref().unwrap(), args.flag_qual, args.flag_shift_counts, args.flag_no_scale, &dist_range, args.flag_only_paired);

        if !args.flag_skip_bed {
            let outfile_bed = if args.flag_no_scale {
                    stem_filename(&bamfile, "_not_scaled.bed", args.flag_bed)
                } else {
                    stem_filename(&bamfile, "_scaled.bed", args.flag_bed)
                };

            match pileup.write_bed(&outfile_bed, args.flag_stranded) {
                Ok(_) => println!("# scale produced {}", &outfile_bed),
                Err(err) => println!("Error producing BED file: {}", err.description()),
            }
        } else {
            println!("# scale skipping BED output");
        }

        if !args.flag_skip_bw {
            let outfile_bw = stem_filename(&bamfile, ".bw", args.flag_bw);

            match pileup.write_bw(&outfile_bw, args.flag_stranded) {
                Ok((f1, f2)) => {
                    if f2.is_some() {
                        println!("# scale produced {}", f1);
                        println!("# scale produced {}", f2.unwrap());
                    } else {
                        println!("# scale produced {}", f1);
                    }
                },
                Err(err) => println!("Error producing BigWig file: {}", err.description()),
            }
        } else {
            println!("# scale skipping BigWig output");
        }
    }
}
<|endoftext|>


<commit_before>
// YAML files as iptables configuration sources

// Rust Core
use std::ascii::AsciiExt;
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::process::exit;

// Crate.io
extern crate clap;
extern crate yaml_rust;

use clap::{Arg, App};
use yaml_rust::{Yaml, YamlLoader};

fn main() {

    let app = App::new("iptables-compose")
        .version("1.1.0")
        .global_version(true)
        .unified_help_message(true)
        .arg_required_else_help(true)
        .about("\nYAML files as iptables configuration sources")
        .arg(Arg::with_name("CONFIG")
             .multiple(false)
             .help("yaml file as iptables configuration source")
             .conflicts_with("license")
             .index(1))
        .arg(Arg::with_name("RESET")
             .short("r")
             .long("reset")
             .help("reset iptables rules\n")
             .requires("CONFIG"))
        .args_from_usage("-l --license 'Prints License'");

    let matches = app.get_matches();

    if let Some(ref f_path) = matches.value_of("CONFIG") {
        // Reset rules if "reset" argument is present
        if matches.is_present("RESET") {
            reset_rules();
        }
        read_yaml(f_path);
    }

    if matches.is_present("license") {
        print_license();
    }

}

fn print_license() {
    let s:String = "
The MIT License (MIT)

Copyright (c) 2015 Samuel <sam@infinitely.io>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the \"Software\"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
".to_string();
    println!("{}", s);
}

fn read_yaml(f_path: &str) {

    // Create a path for yaml configuration file
    let path = Path::new(&f_path);
    let display = path.display();

    // Open file
    let mut file = match File::open(&path) {
        Err(_)   => {
            println!("Failed to open file \"{}\"", display);
            exit(1);
        },
        Ok(file) => file,
    };

    // Create buffer string and read content
    let mut buffer_str = String::new();
    match file.read_to_string(&mut buffer_str) {
        Err(_)  => {
            println!("Failed to read file \"{}\"", display);
            exit(1);
        },
        Ok(str) => str,
    };

    // Parse yaml file
    let yaml = match YamlLoader::load_from_str(&buffer_str) {
        Err(_)  => {
            println!("Failed to parse yaml data: \"{}\"", display);
            exit(1);
        },
        Ok(yaml) => yaml,
    };

    // Check if yaml file has no data, exits program if no data
    if yaml.is_empty() {
        println!("\"{}\" has no configuration(s)", display);
        exit(1);
    };

    // Validate configuration template format
    match yaml[0].as_hash() {
        Some(doc) => doc,
        None => {
            println!("Configuration template format is invalid: {}", display);
            exit(1);
        }
    };

    // Parse configuration from YAML configuration file
    parse_yaml(&yaml[0]);

}

fn parse_yaml(doc: &Yaml) {
    // Read all rules from yaml
    match doc {
        // Parse if template data is a hash object
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                match k.as_str().unwrap() {
                    // Parse `filter` rules
                    "filter" | "FILTER" => parse_filter(v),
                    // Parse custom section
                    _ => parse_section(v)
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn reset_rules() {
    let mut s:String = "iptables -F".to_string();
    s.push_str("\niptables -X");
    s.push_str("\niptables -t nat -F");
    s.push_str("\niptables -t nat -X");
    s.push_str("\niptables -t mangle -F");
    s.push_str("\niptables -t mangle -X");
    println!("{}", s);
}

fn parse_section(doc: &Yaml) {
    match doc {
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                let k = k.as_str().unwrap();
                match k {
                    "ports" | "PORTS" => parse_ports(v),
                    _ => {
                        println!("Rules \"{}\" is not available", k);
                        exit(1);
                    }
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn parse_ports(doc: &Yaml) {
    match doc {
        &Yaml::Array(ref v) => {
            for x in v {
                parse_port_item(x);
            }
        },
        _ => {
            println!("Configuration \"ports\" format is invalid");
            exit(1);
        }
    }
}

fn parse_port_item(doc: &Yaml) {
    match doc {
        &Yaml::Hash(_) => {
            if doc["port"].is_badvalue() {
                println!("Port is not defined");
            }
            let port = doc["port"].as_i64();
            match port {
                Some(port) if port > -1 => (),
                _ => {
                    if port.unwrap() <= -1 {
                        println!("Port has to be greater or equals to 0");
                    } else {
                        println!("Port is not invalid");
                    }
                    exit(1);
                }
            }
            let mut cmd:String = "iptables".to_string();
            if ! doc["forward"].is_badvalue() {
                cmd.push_str(" -t nat -A PREROUTING");
            } else {
                let direction = doc["type"].as_str().unwrap_or("input");
                match direction {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        cmd.push_str(" -I ");
                        cmd.push_str(&direction.to_ascii_uppercase());
                    },
                    _ => {
                        println!("Direction value is invalid");
                        exit(1);
                    }
                }
            }
            match doc["subnet"] {
                Yaml::Array(ref v) => {
                    cmd.push_str(" -s ");
                    let mut y: i32 = 0;
                    for x in v {
                        if y > 0 {
                            cmd.push_str(",");
                        }
                        cmd.push_str(x.as_str().unwrap_or("0.0.0.0/0"));
                        y += 1;
                    }
                },
                _ => {}
            }
            let protocol = doc["protocol"].as_str().unwrap_or("tcp");
            match protocol {
                "tcp" | "udp" | "TCP" | "UDP" => {
                    cmd.push_str(" -p ");
                    cmd.push_str(protocol);
                    cmd.push_str(" -m ");
                    cmd.push_str(protocol);
                },
                _ => {
                    println!("Protocol value is invalid");
                    exit(1);
                }
            }
            cmd.push_str(" --dport ");
            cmd.push_str(&port.unwrap().to_string());
            if ! doc["forward"].is_badvalue() {
                let forward = doc["forward"].as_i64();
                match forward {
                    Some(forward) if forward > -1 => (),
                    _ => {
                        if forward.unwrap() <= -1 {
                            println!("forward port has to be greater or equals to 0");
                        } else {
                            println!("forward port is not invalid");
                        }
                        exit(1);
                    }
                }
                cmd.push_str(" -j REDIRECT --to-port ");
                cmd.push_str(&forward.unwrap().to_string());
            } else {
                let allow = doc["allow"].as_bool().unwrap_or(true);
                cmd.push_str(" -j ");
                match allow {
                    true => {
                        cmd.push_str("ACCEPT");
                    },
                    _ => {
                        cmd.push_str("DROP");
                    }
                }
            }
            println!("{}", cmd);
        },
        _ => {
            println!("Configuration \"ports\" item format is invalid");
            exit(1);
        }
    }
}

fn parse_filter(doc: &Yaml) {
    // Check if filter section exists
    match doc {
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                let k = k.as_str().unwrap();
                let v = v.as_str().unwrap();
                match k {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        match v {
                            "drop" | "reject" | "accept" | "DROP" | "REJECT" | "ACCEPT" => println!("iptables -P {} {}", k.to_ascii_uppercase(), v.to_ascii_uppercase()),
                            _ => {
                                println!("Rules \"{}\" only accept options of \"drop\",\"reject\" or \"accept\"", k);
                                exit(1);
                            }
                        }
                    },
                    _ => println!("iptables -P {}", k.to_ascii_uppercase())
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}
<commit_msg>exit when error<commit_after>
// YAML files as iptables configuration sources

// Rust Core
use std::ascii::AsciiExt;
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::process::exit;

// Crate.io
extern crate clap;
extern crate yaml_rust;

use clap::{Arg, App};
use yaml_rust::{Yaml, YamlLoader};

fn main() {

    let app = App::new("iptables-compose")
        .version("1.1.0")
        .global_version(true)
        .unified_help_message(true)
        .arg_required_else_help(true)
        .about("\nYAML files as iptables configuration sources")
        .arg(Arg::with_name("CONFIG")
             .multiple(false)
             .help("yaml file as iptables configuration source")
             .conflicts_with("license")
             .index(1))
        .arg(Arg::with_name("RESET")
             .short("r")
             .long("reset")
             .help("reset iptables rules\n")
             .requires("CONFIG"))
        .args_from_usage("-l --license 'Prints License'");

    let matches = app.get_matches();

    if let Some(ref f_path) = matches.value_of("CONFIG") {
        // Reset rules if "reset" argument is present
        if matches.is_present("RESET") {
            reset_rules();
        }
        read_yaml(f_path);
    }

    if matches.is_present("license") {
        print_license();
    }

}

fn print_license() {
    let s:String = "
The MIT License (MIT)

Copyright (c) 2015 Samuel <sam@infinitely.io>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the \"Software\"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
".to_string();
    println!("{}", s);
}

fn read_yaml(f_path: &str) {

    // Create a path for yaml configuration file
    let path = Path::new(&f_path);
    let display = path.display();

    // Open file
    let mut file = match File::open(&path) {
        Err(_)   => {
            println!("Failed to open file \"{}\"", display);
            exit(1);
        },
        Ok(file) => file,
    };

    // Create buffer string and read content
    let mut buffer_str = String::new();
    match file.read_to_string(&mut buffer_str) {
        Err(_)  => {
            println!("Failed to read file \"{}\"", display);
            exit(1);
        },
        Ok(str) => str,
    };

    // Parse yaml file
    let yaml = match YamlLoader::load_from_str(&buffer_str) {
        Err(_)  => {
            println!("Failed to parse yaml data: \"{}\"", display);
            exit(1);
        },
        Ok(yaml) => yaml,
    };

    // Check if yaml file has no data, exits program if no data
    if yaml.is_empty() {
        println!("\"{}\" has no configuration(s)", display);
        exit(1);
    };

    // Validate configuration template format
    match yaml[0].as_hash() {
        Some(doc) => doc,
        None => {
            println!("Configuration template format is invalid: {}", display);
            exit(1);
        }
    };

    // Parse configuration from YAML configuration file
    parse_yaml(&yaml[0]);

}

fn parse_yaml(doc: &Yaml) {
    // Read all rules from yaml
    match doc {
        // Parse if template data is a hash object
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                match k.as_str().unwrap() {
                    // Parse `filter` rules
                    "filter" | "FILTER" => parse_filter(v),
                    // Parse custom section
                    _ => parse_section(v)
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn reset_rules() {
    let mut s:String = "iptables -F".to_string();
    s.push_str("\niptables -X");
    s.push_str("\niptables -t nat -F");
    s.push_str("\niptables -t nat -X");
    s.push_str("\niptables -t mangle -F");
    s.push_str("\niptables -t mangle -X");
    println!("{}", s);
}

fn parse_section(doc: &Yaml) {
    match doc {
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                let k = k.as_str().unwrap();
                match k {
                    "ports" | "PORTS" => parse_ports(v),
                    _ => {
                        println!("Rules \"{}\" is not available", k);
                        exit(1);
                    }
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}

fn parse_ports(doc: &Yaml) {
    match doc {
        &Yaml::Array(ref v) => {
            for x in v {
                parse_port_item(x);
            }
        },
        _ => {
            println!("Configuration \"ports\" format is invalid");
            exit(1);
        }
    }
}

fn parse_port_item(doc: &Yaml) {
    match doc {
        &Yaml::Hash(_) => {
            if doc["port"].is_badvalue() {
                println!("Port is not defined");
                exit(1);
            }
            let port = doc["port"].as_i64();
            match port {
                Some(port) if port > -1 => (),
                _ => {
                    if port.unwrap() <= -1 {
                        println!("Port has to be greater or equals to 0");
                    } else {
                        println!("Port is not invalid");
                    }
                    exit(1);
                }
            }
            let mut cmd:String = "iptables".to_string();
            if ! doc["forward"].is_badvalue() {
                cmd.push_str(" -t nat -A PREROUTING");
            } else {
                let direction = doc["type"].as_str().unwrap_or("input");
                match direction {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        cmd.push_str(" -I ");
                        cmd.push_str(&direction.to_ascii_uppercase());
                    },
                    _ => {
                        println!("Direction value is invalid");
                        exit(1);
                    }
                }
            }
            match doc["subnet"] {
                Yaml::Array(ref v) => {
                    cmd.push_str(" -s ");
                    let mut y: i32 = 0;
                    for x in v {
                        if y > 0 {
                            cmd.push_str(",");
                        }
                        cmd.push_str(x.as_str().unwrap_or("0.0.0.0/0"));
                        y += 1;
                    }
                },
                _ => {}
            }
            let protocol = doc["protocol"].as_str().unwrap_or("tcp");
            match protocol {
                "tcp" | "udp" | "TCP" | "UDP" => {
                    cmd.push_str(" -p ");
                    cmd.push_str(protocol);
                    cmd.push_str(" -m ");
                    cmd.push_str(protocol);
                },
                _ => {
                    println!("Protocol value is invalid");
                    exit(1);
                }
            }
            cmd.push_str(" --dport ");
            cmd.push_str(&port.unwrap().to_string());
            if ! doc["forward"].is_badvalue() {
                let forward = doc["forward"].as_i64();
                match forward {
                    Some(forward) if forward > -1 => (),
                    _ => {
                        if forward.unwrap() <= -1 {
                            println!("forward port has to be greater or equals to 0");
                        } else {
                            println!("forward port is not invalid");
                        }
                        exit(1);
                    }
                }
                cmd.push_str(" -j REDIRECT --to-port ");
                cmd.push_str(&forward.unwrap().to_string());
            } else {
                let allow = doc["allow"].as_bool().unwrap_or(true);
                cmd.push_str(" -j ");
                match allow {
                    true => {
                        cmd.push_str("ACCEPT");
                    },
                    _ => {
                        cmd.push_str("DROP");
                    }
                }
            }
            println!("{}", cmd);
        },
        _ => {
            println!("Configuration \"ports\" item format is invalid");
            exit(1);
        }
    }
}

fn parse_filter(doc: &Yaml) {
    // Check if filter section exists
    match doc {
        &Yaml::Hash(ref h) => {
            for (k, v) in h {
                let k = k.as_str().unwrap();
                let v = v.as_str().unwrap();
                match k {
                    "input" | "output" | "forward" | "INPUT" | "OUTPUT" | "FORWARD" => {
                        match v {
                            "drop" | "reject" | "accept" | "DROP" | "REJECT" | "ACCEPT" => println!("iptables -P {} {}", k.to_ascii_uppercase(), v.to_ascii_uppercase()),
                            _ => {
                                println!("Rules \"{}\" only accept options of \"drop\",\"reject\" or \"accept\"", k);
                                exit(1);
                            }
                        }
                    },
                    _ => println!("iptables -P {}", k.to_ascii_uppercase())
                }
            }
        },
        _ => {
            println!("Configuration template format is invalid");
            exit(1);
        }
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;
use std::net::TcpStream;
use std::fs::File;
use std::env;

extern crate chrono;
use chrono::*;

#[macro_use]
extern crate clap;

extern crate env_logger;
#[macro_use]
extern crate log;

extern crate stats;

#[macro_use]
extern crate prometheus;
extern crate hyper;

use hyper::header::ContentType;
use hyper::server;
use hyper::mime::Mime;

use prometheus::{Gauge, Encoder, TextEncoder};

mod analyzer;
mod args;
mod filter;
mod log_parser;
use log_parser::log_events::*;
mod render;
mod request_response_matcher;
use request_response_matcher::*;

fn main() {
    env_logger::init().expect("Failed to initialize logging.");

    let args = args::parse_args(env::args()).expect("Failed to parse arguments.");

    fn run(args: &args::RequestLogAnalyzerArgs) -> Option<analyzer::RequestLogAnalyzerResult> {
        let input: Box<io::Read> = match args.filename.as_ref() {
            "-" => Box::new(io::stdin()),
            _ => Box::new(File::open(&args.filename).expect("Failed to open file.")),
        };

        let timings = extract_timings(input, &args.conditions);
        let result = analyzer::analyze(&timings);

        result
    };

    let result = run(&args);
    run(&args);

    let mut stream;
    let mut renderer: Box<render::Renderer>;

    renderer = match args.graphite_server {
        Some(ref graphite_server) => {
            stream = TcpStream::connect((graphite_server.as_ref(), args.graphite_port.unwrap()))
                .expect("Could not connect to the Graphite server");

            Box::new(render::GraphiteRenderer::new(UTC::now(),
                                                   args.graphite_prefix.clone(),
                                                   &mut stream))
        }
        None => Box::new(render::TerminalRenderer::new()),
    };

    match result {
        Some(result) => {
            renderer.render(result);
        }
        None => warn!("No matching log lines in file."),
    }

    let http_body_gauge: Gauge = register_gauge!(opts!("example_http_response_size_bytes",
                                                       "The HTTP response sizes in bytes.",
                                                       labels!{"handler" => "all",}))
        .unwrap();

    let encoder = TextEncoder::new();
    let addr = "127.0.0.1:9898";
    println!("listening addr {:?}", addr);
    hyper::server::Server::http(addr)
        .unwrap()
        .handle(move |_: hyper::server::Request, mut res: hyper::server::Response| {

            let result = run(&args);

            http_body_gauge.set(result.unwrap().count as f64);

            let metric_familys = prometheus::gather();
            let mut buffer = vec![];

            encoder.encode(&metric_familys, &mut buffer).unwrap();

            res.headers_mut()
                .set(ContentType(encoder.format_type().parse::<Mime>().unwrap()));
            res.send(&buffer).unwrap();
        })
        .unwrap();
}

fn extract_timings(input: Box<io::Read>, conditions: &filter::FilterConditions) -> Vec<i64> {
    let reader = io::BufReader::new(input);

    let mut requests: Vec<Request> = Vec::new();
    let mut responses: Vec<Response> = Vec::new();
    let mut timings: Vec<i64> = Vec::new();

    for line in reader.lines() {
        let parsed_line = log_parser::parse_line(&line.expect("Failed to read line."));

        match parsed_line {
            Ok(event) => {
                match event {
                    LogEvent::Request(request) => requests.push(request),
                    LogEvent::Response(response) => responses.push(response),
                }

                let pairs = extract_matching_request_response_pairs(&mut requests, &mut responses);

                let mut new_timings: Vec<i64> = pairs.iter()
                    .filter(|pair| filter::matches_filter(&pair, conditions))
                    .map(|pair| pair.response.response_time.num_milliseconds())
                    .collect();

                timings.append(&mut new_timings);
            }
            Err(err) => warn!("{}", err),
        }
    }
    timings
}

#[test]
fn test_extract_timings() {
    let conditions = filter::FilterConditions {
        include_terms: None,
        exclude_terms: None,
        latest_time: None,
    };

    let timings = extract_timings(Box::new(File::open("src/test/simple-1.log").unwrap()),
                                  &conditions);

    assert_eq!(timings, vec![7, 10]);
}

#[test]
fn test_extract_timings_ignore_broken_lines() {
    let conditions = filter::FilterConditions {
        include_terms: None,
        exclude_terms: None,
        latest_time: None,
    };

    let timings = extract_timings(Box::new(File::open("src/test/broken.log").unwrap()),
                                  &conditions);

    assert_eq!(timings, vec![7]);
}
<commit_msg>Render real result with Prometheus renderer<commit_after>use std::io;
use std::io::prelude::*;
use std::net::TcpStream;
use std::fs::File;
use std::env;

extern crate chrono;
use chrono::*;

#[macro_use]
extern crate clap;

extern crate env_logger;
#[macro_use]
extern crate log;

extern crate stats;

extern crate prometheus;
extern crate hyper;

use hyper::header::ContentType;
use hyper::server;
use hyper::mime::Mime;

mod analyzer;
mod args;
mod filter;
mod log_parser;
use log_parser::log_events::*;
mod render;
use render::Renderer;
mod request_response_matcher;
use request_response_matcher::*;

fn main() {
    env_logger::init().expect("Failed to initialize logging.");

    let args = args::parse_args(env::args()).expect("Failed to parse arguments.");

    fn run(args: &args::RequestLogAnalyzerArgs) -> Option<analyzer::RequestLogAnalyzerResult> {
        let input: Box<io::Read> = match args.filename.as_ref() {
            "-" => Box::new(io::stdin()),
            _ => Box::new(File::open(&args.filename).expect("Failed to open file.")),
        };

        let timings = extract_timings(input, &args.conditions);
        let result = analyzer::analyze(&timings);

        result
    };

    let result = run(&args);
    run(&args);

    let mut stream;
    let mut renderer: Box<render::Renderer>;

    renderer = match args.graphite_server {
        Some(ref graphite_server) => {
            stream = TcpStream::connect((graphite_server.as_ref(), args.graphite_port.unwrap()))
                .expect("Could not connect to the Graphite server");

            Box::new(render::GraphiteRenderer::new(UTC::now(),
                                                   args.graphite_prefix.clone(),
                                                   &mut stream))
        }
        None => Box::new(render::TerminalRenderer::new()),
    };

    match result {
        Some(result) => {
            renderer.render(result);
        }
        None => warn!("No matching log lines in file."),
    }

    let addr = "127.0.0.1:9898";
    println!("listening addr {:?}", addr);
    hyper::server::Server::http(addr)
        .unwrap()
        .handle(move |_: hyper::server::Request, mut res: hyper::server::Response| {

            let result = run(&args);

            let mut buffer = vec![];

            {
                let mut renderer = render::prometheus::PrometheusRenderer::new(&mut buffer);
                renderer.render(result.unwrap());
            }
            // res.headers_mut()
            //     .set(ContentType(encoder.format_type().parse::<Mime>().unwrap()));
            res.send(&buffer).unwrap();
        })
        .unwrap();
}

fn extract_timings(input: Box<io::Read>, conditions: &filter::FilterConditions) -> Vec<i64> {
    let reader = io::BufReader::new(input);

    let mut requests: Vec<Request> = Vec::new();
    let mut responses: Vec<Response> = Vec::new();
    let mut timings: Vec<i64> = Vec::new();

    for line in reader.lines() {
        let parsed_line = log_parser::parse_line(&line.expect("Failed to read line."));

        match parsed_line {
            Ok(event) => {
                match event {
                    LogEvent::Request(request) => requests.push(request),
                    LogEvent::Response(response) => responses.push(response),
                }

                let pairs = extract_matching_request_response_pairs(&mut requests, &mut responses);

                let mut new_timings: Vec<i64> = pairs.iter()
                    .filter(|pair| filter::matches_filter(&pair, conditions))
                    .map(|pair| pair.response.response_time.num_milliseconds())
                    .collect();

                timings.append(&mut new_timings);
            }
            Err(err) => warn!("{}", err),
        }
    }
    timings
}

#[test]
fn test_extract_timings() {
    let conditions = filter::FilterConditions {
        include_terms: None,
        exclude_terms: None,
        latest_time: None,
    };

    let timings = extract_timings(Box::new(File::open("src/test/simple-1.log").unwrap()),
                                  &conditions);

    assert_eq!(timings, vec![7, 10]);
}

#[test]
fn test_extract_timings_ignore_broken_lines() {
    let conditions = filter::FilterConditions {
        include_terms: None,
        exclude_terms: None,
        latest_time: None,
    };

    let timings = extract_timings(Box::new(File::open("src/test/broken.log").unwrap()),
                                  &conditions);

    assert_eq!(timings, vec![7]);
}
<|endoftext|>


<commit_before>use gl;
use opencl;
use opencl::mem::CLBuffer;
use std::io::timer;
use std::mem;
use std::time::duration::Duration;
use sdl2;
use sdl2::event::Event;
use stopwatch::TimerSet;
use yaglw::gl_context::{GLContext, GLContextExistence};
use yaglw::shader::Shader;
use yaglw::vertex_buffer::{GLArray, GLBuffer, GLType, VertexAttribData, DrawMode};

static WINDOW_WIDTH: u32 = 800;
static WINDOW_HEIGHT: u32 = 600;

#[repr(C)]
struct RGB {
  pub r: f32,
  pub g: f32,
  pub b: f32,
}

pub fn main() {
  let timers = TimerSet::new();

  let window = make_window();

  let _sdl_gl_context = window.gl_create_context().unwrap();

  // Load the OpenGL function pointers.
  gl::load_with(|s| unsafe {
    mem::transmute(sdl2::video::gl_get_proc_address(s))
  });

  let (gl, mut gl_context) = unsafe {
    GLContext::new()
  };

  match gl_context.get_error() {
    gl::NO_ERROR => {},
    err => {
      println!("OpenGL error 0x{:x} in setup", err);
      return;
    },
  }

  let shader = make_shader(&gl);
  shader.use_shader(&mut gl_context);

  let vao = timers.time("make_picture", || {
    make_picture(&gl, &mut gl_context, &shader)
  });

  while !quit_event() {
    timers.time("draw", || {
      gl_context.clear_buffer();
      vao.draw(&mut gl_context);
      // swap buffers
      window.gl_swap_window();
    });

    timer::sleep(Duration::milliseconds(10));
  }

  timers.print();
}

fn make_shader<'a>(
  gl: &'a GLContextExistence,
) -> Shader<'a> {
  let vertex_shader: String = format!("
    #version 330 core

    const int W = {};
    const int H = {};

    in vec3 color;
    out vec4 v_color;

    void main() {{
      v_color = vec4(color, 1);
      gl_Position =
        vec4(
          float(gl_VertexID % W) / W * 2 - 1,
          float(gl_VertexID / W) / H * 2 - 1,
          0, 1
        );
    }}
  ", WINDOW_WIDTH, WINDOW_HEIGHT);

  let fragment_shader: String = "
    #version 330 core

    in vec4 v_color;

    void main() {
      gl_FragColor = v_color;
    }
  ".to_string();

  let components = vec!(
    ((vertex_shader, gl::VERTEX_SHADER)),
    ((fragment_shader, gl::FRAGMENT_SHADER)),
  );

  Shader::new(gl, components.into_iter())
}

fn make_picture<'a>(
  gl: &'a GLContextExistence,
  gl_context: &mut GLContext,
  shader: &Shader<'a>,
) -> GLArray<'a, RGB> {
  let dat = do_opencl();

  let mut vbo = GLBuffer::new(gl, gl_context, dat.len());
  vbo.push(gl_context, dat.as_slice());

  let attribs = [
    VertexAttribData {
      name: "color",
      size: 3,
      unit: GLType::Float,
    },
  ];

  GLArray::new(
    gl,
    gl_context,
    shader,
    &attribs,
    DrawMode::Points,
    vbo,
  )
}

fn do_opencl() -> Vec<RGB> {
  let (device, ctx, queue) = opencl::util::create_compute_context().unwrap();

  let len = WINDOW_WIDTH * WINDOW_HEIGHT;
  let len = len as usize;

  let output_buffer: CLBuffer<RGB> = ctx.create_buffer(len, opencl::cl::CL_MEM_WRITE_ONLY);

  let program = {
    let ker = format!("
      __kernel void color(__global float *output) {{
        int W = {};
        int H = {};
        int i = get_global_id(0);
        float x = i % W;
        float y = i / W;
        i = 3 * i;
        output[i] = 0;
        output[i + 1] = x / W;
        output[i + 2] = y / H;
      }}
    ", WINDOW_WIDTH, WINDOW_HEIGHT);
    ctx.create_program_from_source(ker.as_slice())
  };
  program.build(&device).ok().expect("Couldn't build program.");

  let kernel = program.create_kernel("color");
  // This is sketchy; we "implicitly cast" output_buffer from a CLBuffer<RGB> to a CLBuffer<f32>.
  kernel.set_arg(0, &output_buffer);

  let event = queue.enqueue_async_kernel(&kernel, len, None, ());

  queue.get(&output_buffer, &event)
}

fn make_window() -> sdl2::video::Window {
  sdl2::init(sdl2::INIT_EVERYTHING);

  sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextMajorVersion, 3);
  sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextMinorVersion, 3);
  sdl2::video::gl_set_attribute(
    sdl2::video::GLAttr::GLContextProfileMask,
    sdl2::video::GLProfile::GLCoreProfile as isize
  );

  sdl2::video::Window::new(
    "OpenCL",
    sdl2::video::WindowPos::PosCentered,
    sdl2::video::WindowPos::PosCentered,
    800,
    600,
    sdl2::video::OPENGL,
  ).unwrap()
}

fn quit_event() -> bool {
  loop {
    match sdl2::event::poll_event() {
      Event::None => {
        return false;
      },
      Event::Quit(_) => {
        return true;
      }
      Event::AppTerminating(_) => {
        return true;
      }
      _ => {},
    }
  }
}
<commit_msg>Hello old friend<commit_after>use gl;
use opencl;
use opencl::mem::CLBuffer;
use std::io::timer;
use std::mem;
use std::time::duration::Duration;
use sdl2;
use sdl2::event::Event;
use stopwatch::TimerSet;
use yaglw::gl_context::{GLContext, GLContextExistence};
use yaglw::shader::Shader;
use yaglw::vertex_buffer::{GLArray, GLBuffer, GLType, VertexAttribData, DrawMode};

static WINDOW_WIDTH: u32 = 800;
static WINDOW_HEIGHT: u32 = 600;

#[repr(C)]
struct RGB {
  pub r: f32,
  pub g: f32,
  pub b: f32,
}

pub fn main() {
  let timers = TimerSet::new();

  let window = make_window();

  let _sdl_gl_context = window.gl_create_context().unwrap();

  // Load the OpenGL function pointers.
  gl::load_with(|s| unsafe {
    mem::transmute(sdl2::video::gl_get_proc_address(s))
  });

  let (gl, mut gl_context) = unsafe {
    GLContext::new()
  };

  match gl_context.get_error() {
    gl::NO_ERROR => {},
    err => {
      println!("OpenGL error 0x{:x} in setup", err);
      return;
    },
  }

  let shader = make_shader(&gl);
  shader.use_shader(&mut gl_context);

  let vao = timers.time("make_picture", || {
    make_picture(&gl, &mut gl_context, &shader)
  });

  while !quit_event() {
    timers.time("draw", || {
      gl_context.clear_buffer();
      vao.draw(&mut gl_context);
      // swap buffers
      window.gl_swap_window();
    });

    timer::sleep(Duration::milliseconds(10));
  }

  timers.print();
}

fn make_shader<'a>(
  gl: &'a GLContextExistence,
) -> Shader<'a> {
  let vertex_shader: String = format!("
    #version 330 core

    const int W = {};
    const int H = {};

    in vec3 color;
    out vec4 v_color;

    void main() {{
      v_color = vec4(color, 1);
      gl_Position =
        vec4(
          float(gl_VertexID % W) / W * 2 - 1,
          float(gl_VertexID / W) / H * 2 - 1,
          0, 1
        );
    }}
  ", WINDOW_WIDTH, WINDOW_HEIGHT);

  let fragment_shader: String = "
    #version 330 core

    in vec4 v_color;

    void main() {
      gl_FragColor = v_color;
    }
  ".to_string();

  let components = vec!(
    ((vertex_shader, gl::VERTEX_SHADER)),
    ((fragment_shader, gl::FRAGMENT_SHADER)),
  );

  Shader::new(gl, components.into_iter())
}

fn make_picture<'a>(
  gl: &'a GLContextExistence,
  gl_context: &mut GLContext,
  shader: &Shader<'a>,
) -> GLArray<'a, RGB> {
  let dat = do_opencl();

  let mut vbo = GLBuffer::new(gl, gl_context, dat.len());
  vbo.push(gl_context, dat.as_slice());

  let attribs = [
    VertexAttribData {
      name: "color",
      size: 3,
      unit: GLType::Float,
    },
  ];

  GLArray::new(
    gl,
    gl_context,
    shader,
    &attribs,
    DrawMode::Points,
    vbo,
  )
}

fn do_opencl() -> Vec<RGB> {
  let (device, ctx, queue) = opencl::util::create_compute_context().unwrap();

  let len = WINDOW_WIDTH as usize * WINDOW_HEIGHT as usize;

  let output_buffer: CLBuffer<RGB> = ctx.create_buffer(len, opencl::cl::CL_MEM_WRITE_ONLY);

  let program = {
    let ker = format!("
        __kernel void color(__global float * output) {{
          int W = {};
          int H = {};

          int maxIt = 128;
          float R = 100;

          int i = get_global_id(0);

          float l = -2;
          float r = 2;
          float b = -2;
          float t = 2;

          float c_x = i % W;
          float c_y = i / W;
          c_x = (c_x / W) * (r - l) + l;
          c_y = (c_y / H) * (t - b) + b;

          float x = 0;
          float y = 0;
          int it;
          for (it = 0; it < maxIt; ++it)
          {{
            float x2 = x * x;
            float y2 = y * y;
            if (x2 + y2 > R * R)
              break;
            // Ordering is important here.
            y = 2*x*y + c_y;
            x = x2 - y2 + c_x;
          }}

          i = i * 3;

          if (it < maxIt) {{
            float progress = (float)it / (float)maxIt;
            output[i] = progress;
            output[i + 1] = 1 - 2 * fabs(0.5 - progress);
            output[i + 2] = 0.5 * (1 - progress);
          }} else {{
            output[i] = 0;
            output[i + 1] = 0;
            output[i + 2] = 0;
          }}
        }}
      ", WINDOW_WIDTH, WINDOW_HEIGHT);
    ctx.create_program_from_source(ker.as_slice())
  };
  program.build(&device).unwrap();

  let kernel = program.create_kernel("color");
  // This is sketchy; we "implicitly cast" output_buffer from a CLBuffer<RGB> to a CLBuffer<f32>.
  kernel.set_arg(0, &output_buffer);

  let event = queue.enqueue_async_kernel(&kernel, len, None, ());

  queue.get(&output_buffer, &event)
}

fn make_window() -> sdl2::video::Window {
  sdl2::init(sdl2::INIT_EVERYTHING);

  sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextMajorVersion, 3);
  sdl2::video::gl_set_attribute(sdl2::video::GLAttr::GLContextMinorVersion, 3);
  sdl2::video::gl_set_attribute(
    sdl2::video::GLAttr::GLContextProfileMask,
    sdl2::video::GLProfile::GLCoreProfile as isize
  );

  let window = sdl2::video::Window::new(
    "OpenCL",
    sdl2::video::WindowPos::PosCentered,
    sdl2::video::WindowPos::PosCentered,
    WINDOW_WIDTH as isize,
    WINDOW_HEIGHT as isize,
    sdl2::video::OPENGL,
  ).unwrap();

  window
}

fn quit_event() -> bool {
  loop {
    match sdl2::event::poll_event() {
      Event::None => {
        return false;
      },
      Event::Quit(_) => {
        return true;
      }
      Event::AppTerminating(_) => {
        return true;
      }
      _ => {},
    }
  }
}
<|endoftext|>


<commit_before>extern crate zwreec;
extern crate getopts;
#[macro_use] extern crate log;
extern crate time;
extern crate term;

use std::env;
use std::vec::Vec;
use std::fs::File;
use std::io::Write;

mod utils;
use self::utils::logger;

// shorthand to display program usage
macro_rules! print_usage {
    ($prog:ident, $opts:ident) => {{
        print_stderr!("{}", $opts.usage(&format!("Usage: {} [-hV] [-vq] [-l [LOGFILE]] [-o OUTPUT] INPUT", $prog)));
    }}
}

// Found in:
// http://stackoverflow.com/a/27590832
macro_rules! print_stderr(
    ($($arg:tt)*) => (
        match write!(&mut ::std::io::stderr(), $($arg)* ) {
            Ok(_) => {},
            Err(x) => panic!("Unable to write to stderr: {}", x),
        }
    )
);

fn main() {
    //early init

    // handling commandline parameters
    let args: Vec<String> = env::args().collect();
    let ref program = args[0];
    let mut loggers: Vec<Box<logger::SharedLogger>> = vec![];

    // define options
    let mut opts = getopts::Options::new();
    opts.optflagmulti("v", "verbose", "be more verbose. can be used multiple times.");
    opts.optflag("q", "quiet", "be quiet");
    opts.optflagopt("l", "logfile", "specify log file (default zwreec.log)", "LOGFILE");
    opts.optopt("o", "", "name of the output file", "FILE");
    opts.optflag("h", "help", "display this help and exit");
    opts.optflag("V", "version", "display version");

    let parsed_opts = match opts.parse(&args[1..]) {
        Ok(m)  => { m }
        Err(f) => {
            // parsing error
            // display usage and return
            print_stderr!("{}\n", f.to_string());
            print_usage!(program, opts);
            // TODO: figure out a way to set exit code
            return;
        }
    };

    // examinate options
    if parsed_opts.opt_present("h") {
        // parsed "-h|--help"
        // display usage and return
        print_usage!(program, opts);
        return;
    }

    if parsed_opts.opt_present("V") {
        // parsed "-V|--version"
        // display current version
        println!("{} {}", program, match option_env!("CFG_VERSION") {
            Some(s) => s.to_string(),
            None => format!("{}.{}.{}{}",
                            env!("CARGO_PKG_VERSION_MAJOR"),
                            env!("CARGO_PKG_VERSION_MINOR"),
                            env!("CARGO_PKG_VERSION_PATCH"),
                            option_env!("CARGO_PKG_VERSION_PRE").unwrap_or(""))
        });
        return;
    }

    if parsed_opts.opt_present("v") {
        // parsed "-v|--verbose"
        // set loglevel to verbose
        loggers.push(logger::TermLogger::new(
                match parsed_opts.opt_count("v") {
                    1 => logger::LogLevelFilter::Info,
                    2 => logger::LogLevelFilter::Debug,
                    _ => logger::LogLevelFilter::Trace,
                }));
    } else if parsed_opts.opt_present("q") {
        // parsed "-q|--quiet"
        // set loglevel to error
        loggers.push(logger::TermLogger::new(logger::LogLevelFilter::Error));
    } else {
        // default
        // set loglevel to warn
        loggers.push(logger::TermLogger::new(logger::LogLevelFilter::Warn));
    }

    if parsed_opts.opt_present("l") {
        // parsed "-l|--logfile"
        // sets a logger to output to logfile
        let name = if let Some(n) = parsed_opts.opt_str("l") {
            n
        } else {
            "zwreec.log".to_string()
        };
        loggers.push(logger::FileLogger::new(
                        logger::LogLevelFilter::Trace,
                        File::create(name).unwrap())
            );
    }

    let outfile = if let Some(file) = parsed_opts.opt_str("o") {
        // parsed "-o FILE"
        // set of to filename
        file
    } else {
        // NOTE: string manipulation in rust is stil weird.
        let mut s = String::new();
        s.push_str("a.out");
        s
    };

    let infile = if parsed_opts.free.len() == 1 {
         // check number of 'free' parameter
         // one free parameter is the input file name
         parsed_opts.free[0].clone()
    } else {
        print_stderr!("Input file name missing\n");
        print_usage!(program, opts);
        // TODO: figure out a way to set exit code
        return;
    };

    // activate logger
    // NOTE: I am not sure, if there is a *good* way to do this earlier

    //let _ = SimpleLogger::init(LogLevelFilter::Info);
    let _ = logger::CombinedLogger::init(loggers);

    debug!("parsed command line options");
    info!("main started");
    warn!("WARN");

    // call library
    zwreec::compile(&infile, &outfile);

    // only for testing
    debug!("(1) {}", zwreec::frontend::temp_hello());
    debug!("(2) {}", zwreec::backend::temp_hello());
    debug!("(3) {}", zwreec::file::temp_hello());

    info!("main finished");
}
<commit_msg>Minor code cleanup & spell checking<commit_after>extern crate zwreec;
extern crate getopts;
#[macro_use] extern crate log;
extern crate time;
extern crate term;

use std::env;
use std::vec::Vec;
use std::fs::File;
use std::io::Write;

mod utils;
use self::utils::logger;

// shorthand to display program usage
macro_rules! print_usage {
    ($prog:ident, $opts:ident) => {{
        print_stderr!("{}", $opts.usage(&format!("Usage: {} [-hV] [-vq] [-l [LOGFILE]] [-o OUTPUT] INPUT", $prog)));
    }}
}

// found in:
// http://stackoverflow.com/a/27590832
macro_rules! print_stderr(
    ($($arg:tt)*) => (
        match write!(&mut ::std::io::stderr(), $($arg)* ) {
            Ok(_) => {},
            Err(x) => panic!("Unable to write to stderr: {}", x),
        }
    )
);

fn main() {
    //early init

    // handling command line parameters
    let args: Vec<String> = env::args().collect();
    let ref program = args[0];
    let mut loggers: Vec<Box<logger::SharedLogger>> = vec![];

    // define options
    let mut opts = getopts::Options::new();
    opts.optflagmulti("v", "verbose", "be more verbose. Can be used multiple times.");
    opts.optflag("q", "quiet", "be quiet");
    opts.optflagopt("l", "logfile", "specify log file (default zwreec.log)", "LOGFILE");
    opts.optopt("o", "", "name of the output file", "FILE");
    opts.optflag("h", "help", "display this help and exit");
    opts.optflag("V", "version", "display version");

    let parsed_opts = match opts.parse(&args[1..]) {
        Ok(m)  => { m }
        Err(f) => {
            // parsing error
            // display usage and return
            print_stderr!("{}\n", f.to_string());
            print_usage!(program, opts);
            // TODO: figure out a way to set exit code
            return;
        }
    };

    // examine options
    if parsed_opts.opt_present("h") {
        // parsed "-h|--help"
        // display usage and return
        print_usage!(program, opts);
        return;
    }

    if parsed_opts.opt_present("V") {
        // parsed "-V|--version"
        // display current version
        println!("{} {}", program, match option_env!("CFG_VERSION") {
            Some(s) => s.to_string(),
            None => format!("{}.{}.{}{}",
                            env!("CARGO_PKG_VERSION_MAJOR"),
                            env!("CARGO_PKG_VERSION_MINOR"),
                            env!("CARGO_PKG_VERSION_PATCH"),
                            option_env!("CARGO_PKG_VERSION_PRE").unwrap_or(""))
        });
        return;
    }

    if parsed_opts.opt_present("v") {
        // parsed "-v|--verbose"
        // set log level to verbose
        loggers.push(logger::TermLogger::new(
                match parsed_opts.opt_count("v") {
                    1 => logger::LogLevelFilter::Info,
                    2 => logger::LogLevelFilter::Debug,
                    _ => logger::LogLevelFilter::Trace,
                }));
    } else if parsed_opts.opt_present("q") {
        // parsed "-q|--quiet"
        // set log level to error
        loggers.push(logger::TermLogger::new(logger::LogLevelFilter::Error));
    } else {
        // default
        // set log level to warn
        loggers.push(logger::TermLogger::new(logger::LogLevelFilter::Warn));
    }

    if parsed_opts.opt_present("l") {
        // parsed "-l|--logfile"
        // sets a logger to output to logfile
        let name = if let Some(n) = parsed_opts.opt_str("l") {
            n
        } else {
            "zwreec.log".to_string()
        };
        loggers.push(logger::FileLogger::new(
                        logger::LogLevelFilter::Trace,
                        File::create(name).unwrap())
            );
    }

    let outfile = if let Some(file) = parsed_opts.opt_str("o") {
        // parsed "-o FILE"
        // set of to filename
        file
    } else {
        // NOTE: string manipulation in rust is still weird.
        let mut s = String::new();
        s.push_str("a.out");
        s
    };

    let infile = if parsed_opts.free.len() == 1 {
         // check number of 'free' parameter
         // one free parameter is the input file name
         parsed_opts.free[0].clone()
    } else {
        print_stderr!("Input file name missing\n");
        print_usage!(program, opts);
        // TODO: figure out a way to set exit code
        return;
    };

    // activate logger
    // NOTE: I am not sure, if there is a *good* way to do this earlier

    let _ = logger::CombinedLogger::init(loggers);

    debug!("parsed command line options");
    info!("main started");

    // call library
    zwreec::compile(&infile, &outfile);

    // only for testing
    debug!("(1) {}", zwreec::frontend::temp_hello());
    debug!("(2) {}", zwreec::backend::temp_hello());
    debug!("(3) {}", zwreec::file::temp_hello());

    info!("main finished");
}
<|endoftext|>


<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 cargo_metadata;
extern crate docopt;
extern crate rustc_serialize;

use cargo_metadata::metadata;
use docopt::Docopt;
use std::{env, fs, io, path, process};
use std::io::Write;

const USAGE: &'static str = "
Cargo Fuzz

Usage:
  cargo fuzz --init
  cargo fuzz --fuzz-target TARGET
  cargo fuzz --add TARGET
  cargo fuzz (-h | --help)

Options:
  -h --help              Show this screen.
  --init                 Initialize fuzz folder
  --fuzz-target TARGET   Run with given fuzz target in fuzz/fuzzers
  --add TARGET           Add a new fuzz target
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_init: bool,
    flag_add: Option<String>,
    flag_fuzz_target: Option<String>,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let result = if args.flag_init {
        init_fuzz()
    } else if let Some(target) = args.flag_add {
        add_target(target)
    } else if let Some(target) = args.flag_fuzz_target {
        let result = run_target(target);
        if let Ok(success) = result {
            if success {
                // Can this ever happen?
                Ok(())
            } else {
                println!("Fuzzing found errors!");
                process::exit(-1)
            }
        } else {
            result.map(|_| ())
        }
    } else {
        println!("Invalid arguments. Usage:\n{}", USAGE);
        return;
    };
    if let Err(error) = result {
        println!("Error: {:?}", error);
    }
}

/// Create all the files and folders we need to run
///
/// This will not clone libfuzzer-sys
fn init_fuzz() -> io::Result<()> {
    // todo error handling
    let meta = metadata(None).unwrap();
    let mut p = env::current_dir().unwrap();
    p.push("Cargo.toml");
    let p = p.to_str().unwrap();
    let me = meta.packages.iter().find(|package| package.manifest_path == p).unwrap();

    fs::create_dir("./fuzz")?;
    fs::create_dir("./fuzz/fuzzers")?;

    let mut cargo = fs::File::create(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[package]
name = "{0}-fuzz"
version = "0.0.1"
authors = ["Automatically generated"]

[dependencies.{0}]
path = ".."

# Prevent this from interfering with workspaces
[workspace]
members = ["."]

[[bin]]
name = "fuzzer_script_1"
path = "fuzzers/fuzzer_script_1.rs"
"#, me.name)?;

    let mut ignore = fs::File::create(path::Path::new("./fuzz/.gitignore"))?;

write!(ignore, r#"
target
libfuzzer
"#)?;

    let mut script = fs::File::create(path::Path::new("./fuzz/fuzzers/fuzzer_script_1.rs"))?;
    dummy_target(&mut script)
}

/// Create a dummy fuzz target script at the given path
fn dummy_target(script: &mut fs::File) -> io::Result<()> {
write!(script, r#"
#![no_main]
extern crate fuzzer_sys;

#[export_name="LLVMFuzzerTestOneInput"]
pub extern fn go(data: *const u8, size: isize) -> i32 {{
    // fuzzed code goes here
    0
}}"#)
}

/// Add a new fuzz target script with a given name
fn add_target(target: String) -> io::Result<()> {
    let target_file = format!("fuzz/fuzzers/{}.rs", target);
    let mut script = fs::File::create(path::Path::new(&target_file))?;
    dummy_target(&mut script)?;

    let mut cargo = fs::OpenOptions::new().append(true).open(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[[bin]]
name = "{0}"
path = "fuzzers/{0}.rs"
"#, target)

}

/// Build or rebuild libFuzzer (rebuilds only if the compiler version changed)
///
/// We can't just use libFuzzer as a dependency since libgcc will
/// get compiled with sanitizer support. RUSTFLAGS does not discriminate
/// between build dependencies and regular ones.
///
/// https://github.com/rust-lang/cargo/issues/3739
fn rebuild_libfuzzer() -> io::Result<()> {
    if let Err(_) = env::set_current_dir("./libfuzzer") {
        let mut git = process::Command::new("git");
        let mut cmd = git.arg("clone")
                         .arg("https://github.com/rust-fuzz/libfuzzer-sys.git")
                         .arg("libfuzzer");
        let result = cmd.spawn()?.wait()?;
        if !result.success() {
            return Err(io::Error::new(io::ErrorKind::Other,
                                      "Failed to clone libfuzzer-sys"))
        }
        env::set_current_dir("./libfuzzer")?;
    }
    let mut cmd = process::Command::new("cargo");
    cmd.arg("build")
       .arg("--release")
       .spawn()?
       .wait()?;

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err(io::Error::new(io::ErrorKind::Other,
                                  "Failed to build libfuzzer-sys"))
    }
    env::set_current_dir("..")
}

/// Fuzz a given fuzz target
fn run_target(target: String) -> io::Result<bool> {
    env::set_current_dir("./fuzz")?;
    rebuild_libfuzzer()?;
    let mut cmd = process::Command::new("cargo");
    cmd.arg("rustc")
       .arg("--verbose")
       .arg("--bin")
       .arg(&target)
       .arg("--")
       .arg("-L")
       .arg("libfuzzer/target/release")
       .env("RUSTFLAGS",
            "-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err(io::Error::new(io::ErrorKind::Other, "Failed to build fuzz target"))
    }

    // can't use cargo run since we can't pass -L args to it
    let path = format!("target/debug/{}", target);
    let mut run_cmd = process::Command::new(path);
    let result = run_cmd.spawn()?.wait()?;
    Ok(result.success())
}
<commit_msg>Improve the template fuzz script<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 cargo_metadata;
extern crate docopt;
extern crate rustc_serialize;

use cargo_metadata::metadata;
use docopt::Docopt;
use std::{env, fs, io, path, process};
use std::io::Write;

const USAGE: &'static str = "
Cargo Fuzz

Usage:
  cargo fuzz --init
  cargo fuzz --fuzz-target TARGET
  cargo fuzz --add TARGET
  cargo fuzz (-h | --help)

Options:
  -h --help              Show this screen.
  --init                 Initialize fuzz folder
  --fuzz-target TARGET   Run with given fuzz target in fuzz/fuzzers
  --add TARGET           Add a new fuzz target
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_init: bool,
    flag_add: Option<String>,
    flag_fuzz_target: Option<String>,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let result = if args.flag_init {
        init_fuzz()
    } else if let Some(target) = args.flag_add {
        add_target(target)
    } else if let Some(target) = args.flag_fuzz_target {
        let result = run_target(target);
        if let Ok(success) = result {
            if success {
                // Can this ever happen?
                Ok(())
            } else {
                println!("Fuzzing found errors!");
                process::exit(-1)
            }
        } else {
            result.map(|_| ())
        }
    } else {
        println!("Invalid arguments. Usage:\n{}", USAGE);
        return;
    };
    if let Err(error) = result {
        println!("Error: {:?}", error);
    }
}

/// Create all the files and folders we need to run
///
/// This will not clone libfuzzer-sys
fn init_fuzz() -> io::Result<()> {
    // todo error handling
    let meta = metadata(None).unwrap();
    let mut p = env::current_dir().unwrap();
    p.push("Cargo.toml");
    let p = p.to_str().unwrap();
    let me = meta.packages.iter().find(|package| package.manifest_path == p).unwrap();

    fs::create_dir("./fuzz")?;
    fs::create_dir("./fuzz/fuzzers")?;

    let mut cargo = fs::File::create(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[package]
name = "{0}-fuzz"
version = "0.0.1"
authors = ["Automatically generated"]

[dependencies.{0}]
path = ".."

# Prevent this from interfering with workspaces
[workspace]
members = ["."]

[[bin]]
name = "fuzzer_script_1"
path = "fuzzers/fuzzer_script_1.rs"
"#, me.name)?;

    let mut ignore = fs::File::create(path::Path::new("./fuzz/.gitignore"))?;

write!(ignore, r#"
target
libfuzzer
"#)?;

    let mut script = fs::File::create(path::Path::new("./fuzz/fuzzers/fuzzer_script_1.rs"))?;
    dummy_target(&mut script)
}

/// Create a dummy fuzz target script at the given path
fn dummy_target(script: &mut fs::File) -> io::Result<()> {
write!(script, "{}", r#"
#![no_main]
extern crate fuzzer_sys;

#[export_name="rust_fuzzer_test_input"]
pub extern fn go(data: &[u8]) {
    // fuzzed code goes here
}"#)
}

/// Add a new fuzz target script with a given name
fn add_target(target: String) -> io::Result<()> {
    let target_file = format!("fuzz/fuzzers/{}.rs", target);
    let mut script = fs::File::create(path::Path::new(&target_file))?;
    dummy_target(&mut script)?;

    let mut cargo = fs::OpenOptions::new().append(true).open(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[[bin]]
name = "{0}"
path = "fuzzers/{0}.rs"
"#, target)

}

/// Build or rebuild libFuzzer (rebuilds only if the compiler version changed)
///
/// We can't just use libFuzzer as a dependency since libgcc will
/// get compiled with sanitizer support. RUSTFLAGS does not discriminate
/// between build dependencies and regular ones.
///
/// https://github.com/rust-lang/cargo/issues/3739
fn rebuild_libfuzzer() -> io::Result<()> {
    if let Err(_) = env::set_current_dir("./libfuzzer") {
        let mut git = process::Command::new("git");
        let mut cmd = git.arg("clone")
                         .arg("https://github.com/rust-fuzz/libfuzzer-sys.git")
                         .arg("libfuzzer");
        let result = cmd.spawn()?.wait()?;
        if !result.success() {
            return Err(io::Error::new(io::ErrorKind::Other,
                                      "Failed to clone libfuzzer-sys"))
        }
        env::set_current_dir("./libfuzzer")?;
    }
    let mut cmd = process::Command::new("cargo");
    cmd.arg("build")
       .arg("--release")
       .spawn()?
       .wait()?;

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err(io::Error::new(io::ErrorKind::Other,
                                  "Failed to build libfuzzer-sys"))
    }
    env::set_current_dir("..")
}

/// Fuzz a given fuzz target
fn run_target(target: String) -> io::Result<bool> {
    env::set_current_dir("./fuzz")?;
    rebuild_libfuzzer()?;
    let mut cmd = process::Command::new("cargo");
    cmd.arg("rustc")
       .arg("--verbose")
       .arg("--bin")
       .arg(&target)
       .arg("--")
       .arg("-L")
       .arg("libfuzzer/target/release")
       .env("RUSTFLAGS",
            "-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err(io::Error::new(io::ErrorKind::Other, "Failed to build fuzz target"))
    }

    // can't use cargo run since we can't pass -L args to it
    let path = format!("target/debug/{}", target);
    let mut run_cmd = process::Command::new(path);
    let result = run_cmd.spawn()?.wait()?;
    Ok(result.success())
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Alek Ratzloff
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

extern crate getopts;
extern crate char_iter;

mod trtable;

use std::{env, thread};
use std::sync::mpsc::{channel, Receiver};
use getopts::{Options, Matches};
use trtable::TrTable;

/// Permutes over a single term's possible alternate characters
///
fn permute(index: usize, term: &str, built: String, trtab: &TrTable) {
  // base case
  if index == term.len() {
    println!("{}", built);
    return;
  }

  let c: char = term.chars()
    .nth(index)
    .expect("");

  let list = trtab.get(c);
  if list == None {
    // just add on the next string to the built and keep going
    let mut new_built: String = (&built).to_string();
    new_built.push(c);
    permute(index + 1, &term, new_built, &trtab);
  } else {
    // go through all of the possibilities
    let u_list = list.unwrap();
    for append in u_list {
      let new_built: String = (&built).to_string() + append;
      permute(index + 1, &term, new_built, &trtab);
    }
  }
}

fn all_combos(chars: &Vec<char>, count: usize, built: String) {
  // base case
  if count == 0 {
    println!("{}", &built);
    return;
  }

  for letter in chars {
    let mut new_built: String = (&built).to_string();
    new_built.push(*letter);
    all_combos(&chars, count - 1, new_built);
  }
}

fn parse_opts(args: &Vec<String>) -> (Options, Matches) {
  let mut opts = Options::new();
  opts.optopt("c", "combos", "iterate through all possible printable combinations of N letters. You may specify a comma separated list, e.g. 1,2,4,5",
    "N");
  opts.optflag("", "no-upper", "specific to -c. If this flag is set, all combinations will omit upper case letters.");
  opts.optflag("", "no-lower", "specific to -c. If this flag is set, all combinations will omit lower case letters.");
  opts.optflag("", "no-numbers", "specific to -c. If this flag is set, all combinations will omit numbers.");
  opts.optflag("", "no-symbols", "specific to -c. If this flag is set, all combinations will omit symbols.");
  opts.optopt("p", "permute", "permute over a term with similar letters. Use the -t option to specify a table file",
    "TERM");
  opts.optopt("t", "trtab", "translation table file, in YAML format. Required with the -p option.", "FILENAME");

  opts.optopt("j", "threads", "number of threads to use. Default is 1.", "NUM_THREADS");
  opts.optflag("h", "help", "print this menu");
  let matches = match opts.parse(&args[1..]) {
    Ok(m) => { m },
    Err(f) => { panic!(f.to_string()) },
  };
  (opts, matches)
}

fn main() {
  let args: Vec<String> = env::args().collect();
  let (opts, matches) = parse_opts(&args);

  // print help and exit
  if matches.opt_present("h") {
    let program = match args.get(0) {
      Some(s) => s,
      None => unreachable!(),
    };
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
    return;
  }

  // get the table file if it's there
  let table_file = matches.opt_str("t");
  // get the number of threads; if it doesn't exist, default to 1
  let num_threads = match matches.opt_str("j") {
    Some(n_jobs) => n_jobs.parse::<usize>().unwrap(),
    None => 1
  };

  // permutation
  if let Some(permute_str) = matches.opt_str("p") {
    match table_file {
      Some(trtab_fname) => {
        // load the file
        let trtab = TrTable::load(&trtab_fname)
          .ok()
          .unwrap();
        permute(0, &permute_str, "".to_string(), &trtab);
      },
      None => println!("You must specify a table file with the permute option")
    }
  }
  // all combinations
  if let Some(combos) = matches.opt_str("c") {
    let counts = combos.split(",")
      .map(|n| n.parse::<usize>().unwrap());

    let mut all_chars = Vec::new();
    // upper case letters
    if !matches.opt_present("no-upper") {
      for l in char_iter::new('\x41', '\x5a') {
        all_chars.push(l);
      }
    }
    // lower case letters
    if !matches.opt_present("no-lower") {
      for l in char_iter::new('\x61', '\x7a') {
        all_chars.push(l);
      }
    }
    // numbers
    if !matches.opt_present("no-numbers") {
      for l in char_iter::new('\x30', '\x39') {
        all_chars.push(l);
      }
    }
    // symbols
    if !matches.opt_present("no-symbols") {
      // <space> to /
      for l in char_iter::new('\x20', '\x2F') {
        all_chars.push(l);
      }
      // : to @
      for l in char_iter::new('\x3a', '\x40') {
        all_chars.push(l);
      }
      // [ to `
      for l in char_iter::new('\x5b', '\x60') {
        all_chars.push(l);
      }
      // { to ~
      for l in char_iter::new('\x7b', '\x7e') {
        all_chars.push(l);
      }
    }

    // threads that are currently spawned
    let mut threads: Vec<(thread::JoinHandle<()>, Receiver<i32>)> = Vec::new();

    for n in counts {
      for letter in &all_chars {
        // get whether we have any thread slots open; if not, then
        // block until a thread finishes.
        while threads.len() == num_threads {
          let mut offset = 0;
          for index in 0 .. threads.len() {
            let remove = threads[index - offset].1
              .try_recv()
              .is_ok();

            if remove {
              threads.remove(index - offset);
              offset += 1;
            }
          }
          thread::sleep_ms(1);
        }

        // set up a new channel
        let (tx, rx) = channel();
        // create an intermediate var for the all_chars so we don't have to copy it fifty million times
        let all_chars_imdt = all_chars.clone();
        let l = *letter;
        let child = thread::spawn(move || {
          // shift ownership to the thread
          let all_chars_dup = all_chars_imdt;
          let mut new_built = String::new();

          new_built.push(l);
          all_combos(&all_chars_dup, n - 1, new_built);
          drop(tx.send(0));
        });
        threads.push((child, rx));
      }
    }

    // let the rest of the threads finish up
    for (th, _) in threads {
      drop(th.join());
    }
  }
}
<commit_msg>More code motion/organization<commit_after>/*
 * Copyright (C) 2015 Alek Ratzloff
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

extern crate getopts;
extern crate char_iter;

mod trtable;

use std::{env, thread};
use std::sync::mpsc::{channel, Receiver};
use getopts::{Options, Matches};
use trtable::TrTable;

/// Permutes over a single term's possible alternate characters
///
fn permute(index: usize, term: &str, built: String, trtab: &TrTable) {
  // base case
  if index == term.len() {
    println!("{}", built);
    return;
  }

  let c: char = term.chars()
    .nth(index)
    .expect("");

  let list = trtab.get(c);
  if list == None {
    // just add on the next string to the built and keep going
    let mut new_built: String = (&built).to_string();
    new_built.push(c);
    permute(index + 1, &term, new_built, &trtab);
  } else {
    // go through all of the possibilities
    let u_list = list.unwrap();
    for append in u_list {
      let new_built: String = (&built).to_string() + append;
      permute(index + 1, &term, new_built, &trtab);
    }
  }
}

fn all_combos(chars: &Vec<char>, count: usize, built: String) {
  // base case
  if count == 0 {
    println!("{}", &built);
    return;
  }

  for letter in chars {
    let mut new_built: String = (&built).to_string();
    new_built.push(*letter);
    all_combos(&chars, count - 1, new_built);
  }
}

fn parse_opts(args: &Vec<String>) -> (Options, Matches) {
  let mut opts = Options::new();
  opts.optopt("c", "combos", "iterate through all possible printable combinations of N letters. You may specify a comma separated list, e.g. 1,2,4,5",
    "N");
  opts.optflag("", "no-upper", "specific to -c. If this flag is set, all combinations will omit upper case letters.");
  opts.optflag("", "no-lower", "specific to -c. If this flag is set, all combinations will omit lower case letters.");
  opts.optflag("", "no-numbers", "specific to -c. If this flag is set, all combinations will omit numbers.");
  opts.optflag("", "no-symbols", "specific to -c. If this flag is set, all combinations will omit symbols.");
  opts.optopt("p", "permute", "permute over a term with similar letters. Use the -t option to specify a table file",
    "TERM");
  opts.optopt("t", "trtab", "translation table file, in YAML format. Required with the -p option.", "FILENAME");

  opts.optopt("j", "threads", "number of threads to use. Default is 1.", "NUM_THREADS");
  opts.optflag("h", "help", "print this menu");
  let matches = match opts.parse(&args[1..]) {
    Ok(m) => { m },
    Err(f) => { panic!(f.to_string()) },
  };
  (opts, matches)
}

fn main() {
  let args: Vec<String> = env::args().collect();
  let (opts, matches) = parse_opts(&args);

  // print help and exit
  if matches.opt_present("h") {
    let program = match args.get(0) {
      Some(s) => s,
      None => unreachable!(),
    };
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
    return;
  }

  // get the table file if it's there
  let table_file = matches.opt_str("t");
  // get the number of threads; if it doesn't exist, default to 1
  let num_threads = match matches.opt_str("j") {
    Some(n_jobs) => n_jobs.parse::<usize>().unwrap(),
    None => 1
  };

  // permutation
  if let Some(permute_str) = matches.opt_str("p") {
    match table_file {
      Some(trtab_fname) => {
        // load the file
        let trtab = TrTable::load(&trtab_fname)
          .ok()
          .unwrap();
        permute(0, &permute_str, "".to_string(), &trtab);
      },
      None => println!("You must specify a table file with the permute option")
    }
  }
  // all combinations
  if let Some(combos) = matches.opt_str("c") {
    do_all_combos(matches, &combos, num_threads);
  }
}

fn do_all_combos(matches: Matches, combos: &String, num_threads: usize) {
  let counts = combos.split(",")
    .map(|n| n.parse::<usize>().unwrap());

  let mut all_chars = Vec::new();
  // upper case letters
  if !matches.opt_present("no-upper") {
    for l in char_iter::new('\x41', '\x5a') {
      all_chars.push(l);
    }
  }
  // lower case letters
  if !matches.opt_present("no-lower") {
    for l in char_iter::new('\x61', '\x7a') {
      all_chars.push(l);
    }
  }
  // numbers
  if !matches.opt_present("no-numbers") {
    for l in char_iter::new('\x30', '\x39') {
      all_chars.push(l);
    }
  }
  // symbols
  if !matches.opt_present("no-symbols") {
    // <space> to /
    for l in char_iter::new('\x20', '\x2F') {
      all_chars.push(l);
    }
    // : to @
    for l in char_iter::new('\x3a', '\x40') {
      all_chars.push(l);
    }
    // [ to `
    for l in char_iter::new('\x5b', '\x60') {
      all_chars.push(l);
    }
    // { to ~
    for l in char_iter::new('\x7b', '\x7e') {
      all_chars.push(l);
    }
  }

  // threads that are currently spawned
  let mut threads: Vec<(thread::JoinHandle<()>, Receiver<i32>)> = Vec::new();

  for n in counts {
    for letter in &all_chars {
      // get whether we have any thread slots open; if not, then
      // block until a thread finishes.
      while threads.len() == num_threads {
        let mut offset = 0;
        for index in 0 .. threads.len() {
          let remove = threads[index - offset].1
            .try_recv()
            .is_ok();

          if remove {
            threads.remove(index - offset);
            offset += 1;
          }
        }
        thread::sleep_ms(1);
      }

      // set up a new channel
      let (tx, rx) = channel();
      // create an intermediate var for the all_chars so we don't have to copy it fifty million times
      let all_chars_imdt = all_chars.clone();
      let l = *letter;
      let child = thread::spawn(move || {
        // shift ownership to the thread
        let all_chars_dup = all_chars_imdt;
        let mut new_built = String::new();

        new_built.push(l);
        all_combos(&all_chars_dup, n - 1, new_built);
        drop(tx.send(0));
      });
      threads.push((child, rx));
    }
  }

  // let the rest of the threads finish up
  for (th, _) in threads {
    drop(th.join());
  }
}<|endoftext|>