<commit_before>use clap::Clap;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
        .replace("&", "_");
    if !col.chars().next().unwrap().is_alphabetic() {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: &str,
    delimiter: u8,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::ReaderBuilder::new()
        .flexible(true)
        .delimiter(delimiter)
        .from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            let mut row = row?;
            if row.len() < normalized_cols.len() {
                for _ in 0..normalized_cols.len() - row.len() {
                    row.push_field("");
                }
            }
            stmt.execute(&row).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

#[derive(Clap)]
struct Opts {
    #[clap(long, about = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, about = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, about = "Use '\\t' as the delimiter for the CSV")]
    tab: bool,

    #[clap()]
    paths: Vec<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    if [opts.comma, opts.pipe, opts.tab]
        .iter()
        .filter(|&&v| v)
        .count()
        > 1
    {
        eprintln!("Can't pass more than one of --comma, --pipe, and --tab");
        std::process::exit(1);
    }

    let delim = if opts.comma {
        b','
    } else if opts.pipe {
        b'|'
    } else if opts.tab {
        b'\t'
    } else {
        b','
    };

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if opts.paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", &opts.paths[0], delim)?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in opts.paths.iter().enumerate() {
            let mut col_names =
                _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path, delim)?;
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use super::normalize_col;

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[("abc", "abc"), ("2/6/2000", "c_2_6_2000")] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<commit_msg>Added another test<commit_after>use clap::Clap;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
        .replace("&", "_");
    if !col.chars().next().unwrap().is_alphabetic() {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: &str,
    delimiter: u8,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::ReaderBuilder::new()
        .flexible(true)
        .delimiter(delimiter)
        .from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            let mut row = row?;
            if row.len() < normalized_cols.len() {
                for _ in 0..normalized_cols.len() - row.len() {
                    row.push_field("");
                }
            }
            stmt.execute(&row).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

#[derive(Clap)]
struct Opts {
    #[clap(long, about = "Use ',' as the delimiter for the CSV")]
    comma: bool,
    #[clap(long, about = "Use '|' as the delimiter for the CSV")]
    pipe: bool,
    #[clap(long, about = "Use '\\t' as the delimiter for the CSV")]
    tab: bool,

    #[clap()]
    paths: Vec<String>,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    if [opts.comma, opts.pipe, opts.tab]
        .iter()
        .filter(|&&v| v)
        .count()
        > 1
    {
        eprintln!("Can't pass more than one of --comma, --pipe, and --tab");
        std::process::exit(1);
    }

    let delim = if opts.comma {
        b','
    } else if opts.pipe {
        b'|'
    } else if opts.tab {
        b'\t'
    } else {
        b','
    };

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if opts.paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", &opts.paths[0], delim)?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in opts.paths.iter().enumerate() {
            let mut col_names =
                _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path, delim)?;
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use super::normalize_col;

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[
            ("abc", "abc"),
            ("abc (123)", "abc"),
            ("2/6/2000", "c_2_6_2000"),
        ] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<|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 clap;
extern crate rustc_serialize;
extern crate term;

use cargo_metadata::{metadata, Package};
use clap::{App, Arg, SubCommand, ArgMatches, AppSettings};
use std::{env, error, fs, io, path, process};
use std::io::Write;

mod utils;

fn main() {
    let app = App::new("cargo-fuzz")
        .version(option_env!("CARGO_PKG_VERSION").unwrap_or("0.0.0"))
        .about(option_env!("CARGO_PKG_DESCRIPTION").unwrap_or(""))
        .setting(AppSettings::SubcommandRequired)
        .setting(AppSettings::GlobalVersion)
        .subcommand(SubCommand::with_name("init").about("Initialize the fuzz folder"))
        .subcommand(SubCommand::with_name("run").about("Run the fuzz target in fuzz/fuzzers")
                    .arg(Arg::with_name("TARGET").required(true)))
        .subcommand(SubCommand::with_name("add").about("Add a new fuzz target")
                    .arg(Arg::with_name("TARGET").required(true)))
        .subcommand(SubCommand::with_name("list").about("List all fuzz targets"));
    let args = app.get_matches();

    ::std::process::exit(match args.subcommand() {
        ("init", _) => init_fuzz(),
        ("add", matches) => add_target(matches.expect("arguments present")),
        ("list", _) => list_targets(),
        ("run", matches) => exec_target(matches.expect("arguments present")),
        (s, _) => panic!("unimplemented subcommand {}!", s),
    }.map(|_| 0).unwrap_or_else(|err| {
        writeln!(io::stderr(), "Error: {}", err).expect("failed writing to stderr");
        1
    }));
}

/// Create all the files and folders we need to run
///
/// This will not clone libfuzzer-sys
fn init_fuzz() -> Result<(), Box<error::Error>> {
    let me = get_package();

    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"]
publish = false

[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
corpus
artifacts
"#)?;

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

fn list_targets() -> Result<(), Box<error::Error>> {
    if !path::Path::new("./fuzz").is_dir() {
        return Err("Fuzzing has not been initialized. Run `cargo fuzz init` first.".into());
    }
    env::set_current_dir("./fuzz")?;
    let package = get_package();
    for target in &package.targets {
        println!("{}", target.name);
    }
    Ok(())
}

/// Returns metadata for the Cargo package in the current directory
fn get_package() -> Package {
    // 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();
    meta.packages.into_iter().find(|package| package.manifest_path == p).unwrap()
}

/// If the package contains a library target, generate an `extern crate` line to link to it.
fn link_to_lib(pkg: &Package) -> Option<String> {
    pkg.targets.iter()
               .find(|target| target.kind.iter().any(|k| k == "lib"))
               .map(|target| format!("extern crate {};\n", target.name.replace("-", "_")))
}

/// Create a dummy fuzz target script at the given path
fn dummy_target(script: &mut fs::File, pkg: &Package) -> Result<(), Box<error::Error>> {
write!(script, r#"#![no_main]
extern crate libfuzzer_sys;
{}
#[export_name="rust_fuzzer_test_input"]
pub extern fn go(data: &[u8]) {{
    // fuzzed code goes here
}}"#, link_to_lib(pkg).unwrap_or(String::new())).map_err(|e| e.into())
}

/// Add a new fuzz target script with a given name
fn add_target<'a>(args: &ArgMatches<'a>) -> Result<(), Box<error::Error>> {
    let target: String = args.value_of_os("TARGET").expect("TARGET is required").to_os_string()
        .into_string().map_err(|_| "TARGET must be valid unicode")?;
    let components: &[&path::Path] = &["fuzz".as_ref(), "fuzzers".as_ref(), target.as_ref()];
    let target_file = components.iter().collect::<path::PathBuf>();
    let mut script = fs::File::create(path::Path::new(&target_file))?;
    let me = get_package();
    dummy_target(&mut script, &me)?;

    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).map_err(|e| e.into())

}

/// 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() -> Result<(), Box<error::Error>> {
    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("Failed to clone libfuzzer-sys".into())
        }
        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("Failed to build libfuzzer-sys".into())
    }
    env::set_current_dir("..")
        .map_err(|e| e.into())
}

fn make_dir_if_not_exist(dir: &str) -> Result<(), io::Error> {
    if let Err(k) = fs::create_dir(dir) {
        if k.kind() == io::ErrorKind::AlreadyExists {
            // do nothing
        } else {
            return Err(k);
        }
    }
    Ok(())
}
/// Fuzz a given fuzz target
fn exec_target<'a>(args: &ArgMatches<'a>) -> Result<(), Box<error::Error>> {
    let target: String = args.value_of_os("TARGET").expect("TARGET is required").to_os_string()
        .into_string().map_err(|_| "TARGET must be valid unicode")?;

    env::set_current_dir("./fuzz")?;
    rebuild_libfuzzer()?;
    let mut flags = env::var("RUSTFLAGS").unwrap_or("".into());
    if !flags.is_empty() {
        flags.push(' ');
    }
    flags.push_str("-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");
    let mut cmd = process::Command::new("cargo");
    cmd.arg("rustc")
       .arg("--verbose")
       .arg("--bin")
       .arg(&target)
       .arg("--target")
       .arg("x86_64-unknown-linux-gnu") // won't pass rustflags to build scripts
       .arg("--")
       .arg("-L")
       .arg("libfuzzer/target/release")
       .env("RUSTFLAGS", &flags);

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err("Failed to build fuzz target".into())
    }

    make_dir_if_not_exist("corpus")?;
    make_dir_if_not_exist("artifacts")?;

    // 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);
    run_cmd.arg("-artifact_prefix=artifacts/")
           .arg("corpus") // must be last arg
           .env("ASAN_OPTIONS", "detect_odr_violation=0");
    exec_cmd(&mut run_cmd)?;
    Ok(())
}

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

#[cfg(not(unix))]
fn exec_cmd(cmd: &mut process::Command) -> ::std::io::Result<process::ExitStatus> {
    cmd.status()
}
<commit_msg>Allow passing flag through to the fuzz target<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 clap;
extern crate rustc_serialize;
extern crate term;

use cargo_metadata::{metadata, Package};
use clap::{App, Arg, SubCommand, ArgMatches, AppSettings};
use std::{env, error, fs, io, path, process};
use std::io::Write;

mod utils;

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

    ::std::process::exit(match args.subcommand() {
        ("init", _) => init_fuzz(),
        ("add", matches) => add_target(matches.expect("arguments present")),
        ("list", _) => list_targets(),
        ("run", matches) => exec_target(matches.expect("arguments present")),
        (s, _) => panic!("unimplemented subcommand {}!", s),
    }.map(|_| 0).unwrap_or_else(|err| {
        writeln!(io::stderr(), "Error: {}", err).expect("failed writing to stderr");
        1
    }));
}

/// Create all the files and folders we need to run
///
/// This will not clone libfuzzer-sys
fn init_fuzz() -> Result<(), Box<error::Error>> {
    let me = get_package();

    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"]
publish = false

[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
corpus
artifacts
"#)?;

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

fn list_targets() -> Result<(), Box<error::Error>> {
    if !path::Path::new("./fuzz").is_dir() {
        return Err("Fuzzing has not been initialized. Run `cargo fuzz init` first.".into());
    }
    env::set_current_dir("./fuzz")?;
    let package = get_package();
    for target in &package.targets {
        println!("{}", target.name);
    }
    Ok(())
}

/// Returns metadata for the Cargo package in the current directory
fn get_package() -> Package {
    // 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();
    meta.packages.into_iter().find(|package| package.manifest_path == p).unwrap()
}

/// If the package contains a library target, generate an `extern crate` line to link to it.
fn link_to_lib(pkg: &Package) -> Option<String> {
    pkg.targets.iter()
               .find(|target| target.kind.iter().any(|k| k == "lib"))
               .map(|target| format!("extern crate {};\n", target.name.replace("-", "_")))
}

/// Create a dummy fuzz target script at the given path
fn dummy_target(script: &mut fs::File, pkg: &Package) -> Result<(), Box<error::Error>> {
write!(script, r#"#![no_main]
extern crate libfuzzer_sys;
{}
#[export_name="rust_fuzzer_test_input"]
pub extern fn go(data: &[u8]) {{
    // fuzzed code goes here
}}"#, link_to_lib(pkg).unwrap_or(String::new())).map_err(|e| e.into())
}

/// Add a new fuzz target script with a given name
fn add_target<'a>(args: &ArgMatches<'a>) -> Result<(), Box<error::Error>> {
    let target: String = args.value_of_os("TARGET").expect("TARGET is required").to_os_string()
        .into_string().map_err(|_| "TARGET must be valid unicode")?;
    let components: &[&path::Path] = &["fuzz".as_ref(), "fuzzers".as_ref(), target.as_ref()];
    let target_file = components.iter().collect::<path::PathBuf>();
    let mut script = fs::File::create(path::Path::new(&target_file))?;
    let me = get_package();
    dummy_target(&mut script, &me)?;

    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).map_err(|e| e.into())

}

/// 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() -> Result<(), Box<error::Error>> {
    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("Failed to clone libfuzzer-sys".into())
        }
        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("Failed to build libfuzzer-sys".into())
    }
    env::set_current_dir("..")
        .map_err(|e| e.into())
}

fn make_dir_if_not_exist(dir: &str) -> Result<(), io::Error> {
    if let Err(k) = fs::create_dir(dir) {
        if k.kind() == io::ErrorKind::AlreadyExists {
            // do nothing
        } else {
            return Err(k);
        }
    }
    Ok(())
}
/// Fuzz a given fuzz target
fn exec_target<'a>(args: &ArgMatches<'a>) -> Result<(), Box<error::Error>> {
    let target: String = args.value_of_os("TARGET").expect("TARGET is required").to_os_string()
        .into_string().map_err(|_| "TARGET must be valid unicode")?;
    let exec_args = args.values_of_os("ARGS");

    let target_triple = "x86_64-unknown-linux-gnu";

    env::set_current_dir("./fuzz")?;
    rebuild_libfuzzer()?;
    let mut flags = env::var("RUSTFLAGS").unwrap_or("".into());
    if !flags.is_empty() {
        flags.push(' ');
    }
    flags.push_str("-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");
    let mut cmd = process::Command::new("cargo");
    cmd.arg("rustc")
       .arg("--verbose")
       .arg("--bin")
       .arg(&target)
       .arg("--target")
       .arg(target_triple) // won't pass rustflags to build scripts
       .arg("--")
       .arg("-L")
       .arg("libfuzzer/target/release")
       .env("RUSTFLAGS", &flags);

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err("Failed to build fuzz target".into())
    }

    make_dir_if_not_exist("corpus")?;
    make_dir_if_not_exist("artifacts")?;

    // can't use cargo run since we can't pass -L args to it
    let components: &[&path::Path] = &["target".as_ref(),
                                       target_triple.as_ref(),
                                       "debug".as_ref(),
                                       target.as_ref()];
    let mut run_cmd = process::Command::new(components.iter().collect::<path::PathBuf>());
    run_cmd.arg("-artifact_prefix=artifacts/")
           .env("ASAN_OPTIONS", "detect_odr_violation=0");
    exec_args.map(|args| for arg in args { run_cmd.arg(arg); });
    run_cmd.arg("corpus"); // must be last arg
    exec_cmd(&mut run_cmd)?;
    Ok(())
}

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

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


<commit_before>extern crate byteorder;
use std::io::Cursor;
use byteorder::{LittleEndian, ReadBytesExt};

// Bytecode instruction opcodes. The values of these opcodes should never
// change, to remain compatible with existing bytecode programs.
//
// To make the from_u8 function work, MAX_INST_VARIANT must be kept in sync with
// this enum and every value from 0 to MAX_INST_VARIANT should be assigned to an
// opcode.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
enum Inst {
    Nop   = 0,
    Push  = 1,
    Dup   = 2,
    Pop   = 3,
    Swap  = 4,
    Add   = 5,
    Print = 6,
    Halt  = 7,
    Jump  = 8,
}
const MAX_INST_VARIANT: u8 = 8;

impl Inst {
    fn from_u8(x: u8) -> Option<Self> {
        if x <= MAX_INST_VARIANT {
            // This is safe as long as MAX_INST_VARIANT
            Some(unsafe { std::mem::transmute::<u8, Self>(x) })
        } else {
            None
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum VmError {
    InvalidOpcode,
    UnexpectedProgramEnd, // Hit end of program while reading opcode argument.
    StackOverflow,
    StackUnderflow,
}

fn execute(program: &[u8],
           stack: &mut [i32],
           mut stack_idx: usize) -> Result<usize, VmError> {
    use VmError::*;

    let stack_size = stack.len();
    let mut opcodes = Cursor::new(program);

    while let Ok(opcode) = opcodes.read_u8() {
        let inst = try!(Inst::from_u8(opcode).ok_or(InvalidOpcode));
        match inst {
            Inst::Nop => {},

            Inst::Push => {
                let val = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
                let stack_top = try!(stack.get_mut(stack_idx).ok_or(StackOverflow));
                *stack_top = val;
                stack_idx += 1;
            },

            Inst::Dup => {
                if stack_idx < 1 { return Err(StackUnderflow); }
                if stack_idx >= stack_size { return Err(StackOverflow); }
                unsafe {
                    *stack.get_unchecked_mut(stack_idx) = *stack.get_unchecked(stack_idx - 1);
                }
                stack_idx += 1;
            },

            Inst::Pop => {
                if stack_idx < 1 { return Err(StackUnderflow); }
                stack_idx -= 1;
            },

            Inst::Swap => {
                if stack_idx < 2 { return Err(StackUnderflow); }
                unsafe {
                    let tmp = *stack.get_unchecked(stack_idx - 1);
                    *stack.get_unchecked_mut(stack_idx - 1) = *stack.get_unchecked(stack_idx - 2);
                    *stack.get_unchecked_mut(stack_idx - 2) = tmp;
                }
            },

            Inst::Add => {
                if stack_idx < 2 { return Err(StackUnderflow); }
                unsafe {
                    *stack.get_unchecked_mut(stack_idx - 2) += *stack.get_unchecked(stack_idx - 1);
                }
                stack_idx -= 1;
            },

            Inst::Print => {
                let val = *try!(stack.get(stack_idx - 1).ok_or(StackUnderflow));
                println!("{}", val);
                stack_idx -= 1;
            },

            Inst::Halt => {
                break;
            },

            Inst::Jump => {
                let delta = try!(opcodes.read_i8().or(Err(UnexpectedProgramEnd)));
                let addr = (opcodes.position() as i64 + delta as i64) as u64;
                opcodes.set_position(addr);
            },
        }

        //println!("{:?}", &stack[..stack_idx]);
    }

    Ok(stack_idx)
}

fn main() {
    let program = &[
        Inst::Push as u8, 1, 0, 0, 0,
        Inst::Push as u8, 2, 0, 0, 0,
        Inst::Add as u8,
        Inst::Dup as u8,
        Inst::Print as u8,
        Inst::Halt as u8,
        Inst::Jump as u8, -10i8 as u8,
    ];
    execute(program, &mut [0; 256], 0).unwrap();
}
<commit_msg>Add assembler.<commit_after>extern crate byteorder;
use std::io::Cursor;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

// Bytecode instruction opcodes. The values of these opcodes should never
// change, to remain compatible with existing bytecode programs.
//
// To make the from_u8 function work, MAX_INST_VARIANT must be kept in sync with
// this enum and every value from 0 to MAX_INST_VARIANT should be assigned to an
// opcode.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
enum Inst {
    Nop   = 0,
    Push  = 1,
    Dup   = 2,
    Pop   = 3,
    Swap  = 4,
    Add   = 5,
    Print = 6,
    Halt  = 7,
    Jump  = 8,
}
const MAX_INST_VARIANT: u8 = 8;

impl Inst {
    fn from_u8(x: u8) -> Option<Self> {
        if x <= MAX_INST_VARIANT {
            // This is safe as long as MAX_INST_VARIANT is correct and all the valid enum values
            // are contiguous.
            Some(unsafe { std::mem::transmute::<u8, Self>(x) })
        } else {
            None
        }
    }

    fn num_operands(&self) -> u8 {
        match *self {
            Inst::Push | Inst::Jump => 1,
            _ => 0,
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum VmError {
    InvalidOpcode,
    UnexpectedProgramEnd, // Hit end of program while reading opcode argument.
    StackOverflow,
    StackUnderflow,
}

fn execute(program: &[u8],
           stack: &mut [i32],
           mut stack_idx: usize) -> Result<usize, VmError> {
    use VmError::*;

    let stack_size = stack.len();
    let mut opcodes = Cursor::new(program);

    while let Ok(opcode) = opcodes.read_u8() {
        let inst = try!(Inst::from_u8(opcode).ok_or(InvalidOpcode));
        match inst {
            Inst::Nop => {},

            Inst::Push => {
                let val = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
                let stack_top = try!(stack.get_mut(stack_idx).ok_or(StackOverflow));
                *stack_top = val;
                stack_idx += 1;
            },

            Inst::Dup => {
                if stack_idx < 1 { return Err(StackUnderflow); }
                if stack_idx >= stack_size { return Err(StackOverflow); }
                unsafe {
                    *stack.get_unchecked_mut(stack_idx) = *stack.get_unchecked(stack_idx - 1);
                }
                stack_idx += 1;
            },

            Inst::Pop => {
                if stack_idx < 1 { return Err(StackUnderflow); }
                stack_idx -= 1;
            },

            Inst::Swap => {
                if stack_idx < 2 { return Err(StackUnderflow); }
                unsafe {
                    let tmp = *stack.get_unchecked(stack_idx - 1);
                    *stack.get_unchecked_mut(stack_idx - 1) = *stack.get_unchecked(stack_idx - 2);
                    *stack.get_unchecked_mut(stack_idx - 2) = tmp;
                }
            },

            Inst::Add => {
                if stack_idx < 2 { return Err(StackUnderflow); }
                unsafe {
                    *stack.get_unchecked_mut(stack_idx - 2) += *stack.get_unchecked(stack_idx - 1);
                }
                stack_idx -= 1;
            },

            Inst::Print => {
                let val = *try!(stack.get(stack_idx - 1).ok_or(StackUnderflow));
                println!("{}", val);
                stack_idx -= 1;
            },

            Inst::Halt => {
                break;
            },

            Inst::Jump => {
                let delta = try!(opcodes.read_i32::<LittleEndian>().or(Err(UnexpectedProgramEnd)));
                let operand_size = std::mem::size_of::<i32>() as i64;
                let addr = (opcodes.position() as i64 + delta as i64 - operand_size) as u64;
                opcodes.set_position(addr);
            },
        }

        //println!("{:?}", &stack[..stack_idx]);
    }

    Ok(stack_idx)
}

fn parse_instruction(inst: &str) -> Option<Inst> {
    match inst {
        "nop"   => Some(Inst::Nop),
        "push"  => Some(Inst::Push),
        "dup"   => Some(Inst::Dup),
        "pop"   => Some(Inst::Pop),
        "swap"  => Some(Inst::Swap),
        "add"   => Some(Inst::Add),
        "print" => Some(Inst::Print),
        "halt"  => Some(Inst::Halt),
        "jump"  => Some(Inst::Jump),
        _       => None
    }
}

fn assemble(source: &str) -> Vec<u8> {
    use std::collections::HashMap;

    let mut program: Vec<u8> = Vec::new();
    let mut label_definitions: HashMap<&str, usize> = HashMap::new();
    let mut label_uses: Vec<(&str, usize)> = Vec::new();

    for line in source.lines() {
        let mut tokens = line.split(char::is_whitespace).filter(|s| !s.is_empty());
        let mut first_token = tokens.next();

        // Parse an optional label at the start of the line.
        if let Some(label) = first_token {
            if label.chars().next_back() == Some(':') {
                let label_name = &label[..label.len() - 1];
                if label_definitions.insert(label_name, program.len()).is_some() {
                    panic!("Attempted to redefine label '{}'", label_name);
                }
                first_token = tokens.next()
            }
        }

        // Parse the rest of the line if it's not blank.
        if let Some(opcode) = first_token {
            // Parse the instruction name.
            let inst = parse_instruction(opcode).unwrap_or_else(|| {
                panic!("Unrecognized instruction '{}'", opcode)
            });
            program.push(inst as u8);

            // Parse the operands.
            for _ in 0..inst.num_operands() {
                let operand = tokens.next().unwrap_or_else(|| {
                    panic!("Missing one or more operands after '{}'", opcode)
                });

                if operand.chars().next() == Some('@') {
                    let label_name = &operand[1..];
                    label_uses.push((label_name, program.len()));

                    // Push four zero bytes to be overwritten by the label location later.
                    for _ in 0..4 { program.push(0); }
                } else if let Ok(number) = operand.parse::<i32>() {
                    let operand_index = program.len();
                    for _ in 0..4 { program.push(0); }
                    (&mut program[operand_index..]).write_i32::<LittleEndian>(number).unwrap();
                } else {
                    panic!("Expected label or valid 32-bit signed integer after '{}', not '{}'",
                          opcode, operand);
                }
            }
        }
    }

    // Resolve label references and fill in their values in the bytecode.
    for (label_name, use_index) in label_uses {
        let target_index = *label_definitions.get(label_name).unwrap_or_else(|| {
            panic!("Reference to undefined label '{}'", label_name);
        });
        let delta = target_index as i32 - use_index as i32;
        (&mut program[use_index..]).write_i32::<LittleEndian>(delta).unwrap();
    }

    program
}

fn main() {
    let source = r"
        push 1
        label: push 2
        add
        dup
        print
        jump @label
    ";

    let program = assemble(source);

    execute(&program, &mut [0; 256], 0).unwrap();
}
<|endoftext|>


<commit_before>extern crate mysql;
extern crate toml;
extern crate rustc_serialize;
#[macro_use]
extern crate lazy_static;

mod lib;

use mysql::conn::MyOpts;
use std::default::Default;
use mysql::conn::pool;
use mysql::value::from_value;

use lib::config;
use lib::downloader::DownloadDB;
use lib::downloader::Downloader;
use lib::DownloadError;
use lib::converter::Converter;

use std::fs::{remove_file};

macro_rules! try_option { ($e:expr) => (match $e { Some(x) => x, None => return None }) }

///Move result value out, return with none on err & print
macro_rules! try_reoption { ($e:expr) => (match $e { Ok(x) => x, Err(e) => {println!("{}",e);return None }}) }

static VERSION : &'static str = "0.1"; // String not valid
static SLEEP_MS: u32 = 5000;
static CODE_FAILED: i8 = 3;
static CODE_SUCCESS: i8 = 2;

lazy_static! {
    static ref CONFIG: config::Config = {
        println!("Starting yayd-backend v{}",&VERSION);
        config::init_config() //return
    };

}

//#[repr(i8)] broken, enum not usable as of #10292
// enum QueryCodes {
//     Waiting = 0,
//     InProgress = 1,
//     Finished = 2,
//     Failed = 3,
// }

fn main() {

    let pool = lib::db_connect(mysql_options(), SLEEP_MS);

    let converter = Converter::new(&CONFIG.general.ffmpeg_bin, pool.clone());
    let mut print_pause = true;
    loop {
        if let Some(result) = request_entry(& pool) {
            print_pause = true;
            if result.playlist {
                println!("Playlist not supported atm!");
                //TODO: set playlist entry to errg
            }
            let qid = result.qid.clone();                 //&QueryCodes::InProgress as i32
            lib::set_query_code(&mut pool.get_conn().unwrap(), &1, &result.qid).ok().expect("Failed to set query code!");
            lib::set_query_state(&pool.clone(),&qid, "started", false);

            let succes;
            {
                let mut left_files: Vec<String> = Vec::with_capacity(2);
                succes = match handle_download(result, None, &converter,&mut left_files) {
                    Ok(v) => v,
                    Err(e) => {println!("Error: {:?}", e); false }
                };

                if !left_files.is_empty() {
                    println!("cleaning up files");
                    for i in &left_files {
                        match remove_file(&i) {
                            Ok(_) => (println!("cleaning up {}",i)),
                            Err(e) => println!("unable to remove file '{}' {}",i,e),
                        }
                    }
                }
            }

            let code: i8 = if succes {
                CODE_SUCCESS//QueryCode
            } else {
                CODE_FAILED//QueryCode
            };
            lib::set_query_code(&mut pool.get_conn().unwrap(), &code,&qid).ok().expect("Failed to set query code!");

            let state = if code == 2 {
                "finished"
            } else {
                "failed"
            };
            lib::set_query_state(&pool.clone(),&qid, state, true);

        } else {
            if print_pause { println!("Pausing.."); print_pause = false; }
            std::thread::sleep_ms(SLEEP_MS);
        }
    }
}

///Download handler
///Used by the playlist/file handler to download one file
///Based on the quality it needs to download audio & video separately and convert them together
///In case of a playlist download it depends on the target download folder & if it should be bezipped
///In case of a DMCA we need to download the file via the socket connector,
///which will output a mp3, or if requested, the video but with the highest quality.
///Thus in case of a DMCA we can't pick a quality anymore.
///Also the filename depends on the socket output then.
///
///If it's a non-zipped single file, it's moved after a successful download, converted etc to the
///main folder from which it should be downloadable.
///The original non-ascii & url_encode'd name of the file is stored in the DB
fn handle_download<'a>(downl_db: DownloadDB, folder: Option<String>, converter: &Converter, file_db: &mut Vec<String>) -> Result<bool,DownloadError>{
    //update progress
    let is_zipped = match folder {
        Some(_) => true,
        None => false,
    };
    lib::update_steps(&downl_db.pool.clone(),&downl_db.qid, 1, 0);

    let download = Downloader::new(&downl_db, &CONFIG.general);
    //get filename, check for DMCA
    let mut dmca = false; // "succ." dmca -> file already downloaded

    let name = match download.get_file_name() { // get filename
        Ok(v) => v,
        Err(DownloadError::DMCAError) => { //now request via lib.. // k if( k == Err(DownloadError::DMCAError) )
            println!("DMCA error!");
            match download.lib_request_video() {
                Err(err) => { println!("Offliberty-call error {:?}", err); return Err(err); },
                Ok(v) => { dmca = true; v },
            }
        },
        Err(e) => { // unknown error / video private etc.. abort
            println!("Unknown error: {:?}", e);
            return Err(e);
        },
    };

    let name_http_valid = lib::format_file_name(&name, &download, &downl_db.qid);

    let file_path = format_file_path(&downl_db.qid, folder.clone());
    file_db.push(file_path.clone());
    let save_path = &format_save_path(folder.clone(),&name, &download, &downl_db.qid);

    println!("Filename: {}", name);

    let mut is_splitted_video = false;

    if !dmca {
        //TODO: insert title name for file,
        //copy file to download folder

        is_splitted_video = lib::is_split_container(&downl_db.quality);
        let total_steps = if is_splitted_video {
            4
        } else if download.is_audio() {
            3
        } else {
            2
        };
        lib::update_steps(&downl_db.pool.clone(),&downl_db.qid, 2, total_steps);

        //download video, which is video/audio(m4a)
        try!(download.download_file(&file_path, false));

        if is_splitted_video { // download audio file & convert together
            lib::update_steps(&downl_db.pool.clone(),&downl_db.qid, 3, total_steps);

            let audio_path = format_audio_path(&downl_db.qid, folder.clone());
            file_db.push(audio_path.clone());

            println!("Downloading audio.. {}", audio_path);
            try!(download.download_file(&audio_path, true));

            lib::update_steps(&downl_db.pool.clone(),&downl_db.qid, 4, total_steps);

            match converter.merge_files(&downl_db.qid,&file_path, &audio_path,&save_path) {
                Err(e) => {println!("merge error: {:?}",e); return Err(e);},
                Ok(()) => {},
            }

            try!(remove_file(&audio_path));
            file_db.pop();

        }else{ // we're already done, only need to copy
            if !download.is_audio() {
                try!(lib::move_file(&file_path, &save_path));
            }
        }
    }
    if download.is_audio(){ // if audio-> convert m4a to mp3, which converts directly to downl. dir
        try!(converter.extract_audio(&downl_db.qid, &file_path, &save_path));
        try!(remove_file(&file_path));
    }else{
        if !is_splitted_video {
            try!(lib::move_file(&file_path, &save_path));
        }
    }
    file_db.pop();

    if !is_zipped { // add file to list, except it's for zip-compression later (=folder set)
        lib::add_file_entry(&downl_db.pool.clone(), &downl_db.qid,&name_http_valid, &name);
    }

    //TODO: download file, convert if audio ?!
    Ok(true)
}

///Format save location for file, zip dependent
fn format_file_path(qid: &i64, folder: Option<String>) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}", &CONFIG.general.save_dir, v, qid),
        None => format!("{}/{}", &CONFIG.general.save_dir, qid),
    }
}

///Formats the audio path, based on the qid & optional folders
fn format_audio_path(qid: &i64, folder: Option<String>) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}a", &CONFIG.general.save_dir, v, qid),
        None => format!("{}/{}a", &CONFIG.general.save_dir, qid),
    }
}

///Format save path, dependent on zip option.
fn format_save_path<'a>(folder: Option<String>, name: &str, download: &'a Downloader, qid: &i64) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}", &CONFIG.general.save_dir, v, lib::format_file_name(name,download,qid)),
        None => format!("{}/{}", &CONFIG.general.download_dir, lib::format_file_name(name,download,qid)),
    }
}

///Request an entry from the DB to handle
fn request_entry(pool: & pool::MyPool) -> Option<DownloadDB> {
    let mut conn = try_reoption!(pool.get_conn());
    let mut stmt = try_reoption!(conn.prepare("SELECT queries.qid,url,type,quality FROM querydetails \
                    INNER JOIN queries \
                    ON querydetails.qid = queries.qid \
                    WHERE querydetails.code = 0 \
                    ORDER BY queries.created \
                    LIMIT 1"));
    let mut result = try_reoption!(stmt.execute(&[]));
    let result = try_reoption!(try_option!(result.next())); // result.next().'Some'->value.'unwrap'

    println!("Result: {:?}", result[0]);
    println!("result str: {}", result[1].into_str());
    let download_db = DownloadDB { url: from_value::<String>(&result[1]),
                                    quality: from_value::<i16>(&result[3]),
                                    qid: from_value::<i64>(&result[0]),
                                    codecs: &CONFIG.codecs,
                                    extensions: &CONFIG.extensions,
                                    folder: CONFIG.general.save_dir.clone(),
                                    pool: pool.clone(),
                                    playlist: false, //TEMP
                                    compress: false };
    Some(download_db)
}

///Set dbms connection settings
fn mysql_options() -> MyOpts {
    MyOpts {
        tcp_addr: Some(CONFIG.db.ip.clone()),
        tcp_port: CONFIG.db.port,
        user: Some(CONFIG.db.user.clone()),
        pass: Some(CONFIG.db.password.clone()),
        db_name: Some(CONFIG.db.db.clone()),
        ..Default::default() // set others to default
    }
}<commit_msg>fixed deleting of video source file<commit_after>extern crate mysql;
extern crate toml;
extern crate rustc_serialize;
#[macro_use]
extern crate lazy_static;

mod lib;

use mysql::conn::MyOpts;
use std::default::Default;
use mysql::conn::pool;
use mysql::value::from_value;

use lib::config;
use lib::downloader::DownloadDB;
use lib::downloader::Downloader;
use lib::DownloadError;
use lib::converter::Converter;

use std::fs::{remove_file};

macro_rules! try_option { ($e:expr) => (match $e { Some(x) => x, None => return None }) }

///Move result value out, return with none on err & print
macro_rules! try_reoption { ($e:expr) => (match $e { Ok(x) => x, Err(e) => {println!("{}",e);return None }}) }

static VERSION : &'static str = "0.1"; // String not valid
static SLEEP_MS: u32 = 5000;
static CODE_FAILED: i8 = 3;
static CODE_SUCCESS: i8 = 2;

lazy_static! {
    static ref CONFIG: config::Config = {
        println!("Starting yayd-backend v{}",&VERSION);
        config::init_config() //return
    };

}

//#[repr(i8)] broken, enum not usable as of #10292
// enum QueryCodes {
//     Waiting = 0,
//     InProgress = 1,
//     Finished = 2,
//     Failed = 3,
// }

fn main() {

    let pool = lib::db_connect(mysql_options(), SLEEP_MS);

    let converter = Converter::new(&CONFIG.general.ffmpeg_bin, pool.clone());
    let mut print_pause = true;
    loop {
        if let Some(result) = request_entry(& pool) {
            print_pause = true;
            if result.playlist {
                println!("Playlist not supported atm!");
                //TODO: set playlist entry to errg
            }
            let qid = result.qid.clone();                 //&QueryCodes::InProgress as i32
            lib::set_query_code(&mut pool.get_conn().unwrap(), &1, &result.qid).ok().expect("Failed to set query code!");
            lib::set_query_state(&pool.clone(),&qid, "started", false);

            let succes;
            {
                let mut left_files: Vec<String> = Vec::with_capacity(2);
                succes = match handle_download(result, None, &converter,&mut left_files) {
                    Ok(v) => v,
                    Err(e) => {println!("Error: {:?}", e); false }
                };

                if !left_files.is_empty() {
                    println!("cleaning up files");
                    for i in &left_files {
                        match remove_file(&i) {
                            Ok(_) => (println!("cleaning up {}",i)),
                            Err(e) => println!("unable to remove file '{}' {}",i,e),
                        }
                    }
                }
            }

            let code: i8 = if succes {
                CODE_SUCCESS//QueryCode
            } else {
                CODE_FAILED//QueryCode
            };
            lib::set_query_code(&mut pool.get_conn().unwrap(), &code,&qid).ok().expect("Failed to set query code!");

            let state = if code == 2 {
                "finished"
            } else {
                "failed"
            };
            lib::set_query_state(&pool.clone(),&qid, state, true);

        } else {
            if print_pause { println!("Pausing.."); print_pause = false; }
            std::thread::sleep_ms(SLEEP_MS);
        }
    }
}

///Download handler
///Used by the playlist/file handler to download one file
///Based on the quality it needs to download audio & video separately and convert them together
///In case of a playlist download it depends on the target download folder & if it should be bezipped
///In case of a DMCA we need to download the file via the socket connector,
///which will output a mp3, or if requested, the video but with the highest quality.
///Thus in case of a DMCA we can't pick a quality anymore.
///Also the filename depends on the socket output then.
///
///If it's a non-zipped single file, it's moved after a successful download, converted etc to the
///main folder from which it should be downloadable.
///The original non-ascii & url_encode'd name of the file is stored in the DB
fn handle_download<'a>(downl_db: DownloadDB, folder: Option<String>, converter: &Converter, file_db: &mut Vec<String>) -> Result<bool,DownloadError>{
    //update progress
    let is_zipped = match folder {
        Some(_) => true,
        None => false,
    };
    lib::update_steps(&downl_db.pool.clone(),&downl_db.qid, 1, 0);

    let download = Downloader::new(&downl_db, &CONFIG.general);
    //get filename, check for DMCA
    let mut dmca = false; // "succ." dmca -> file already downloaded

    let name = match download.get_file_name() { // get filename
        Ok(v) => v,
        Err(DownloadError::DMCAError) => { //now request via lib.. // k if( k == Err(DownloadError::DMCAError) )
            println!("DMCA error!");
            match download.lib_request_video() {
                Err(err) => { println!("Offliberty-call error {:?}", err); return Err(err); },
                Ok(v) => { dmca = true; v },
            }
        },
        Err(e) => { // unknown error / video private etc.. abort
            println!("Unknown error: {:?}", e);
            return Err(e);
        },
    };

    let name_http_valid = lib::format_file_name(&name, &download, &downl_db.qid);

    let file_path = format_file_path(&downl_db.qid, folder.clone());
    file_db.push(file_path.clone());
    let save_path = &format_save_path(folder.clone(),&name, &download, &downl_db.qid);

    println!("Filename: {}", name);

    let mut is_splitted_video = false;

    if !dmca {
        //TODO: insert title name for file,
        //copy file to download folder

        is_splitted_video = lib::is_split_container(&downl_db.quality);
        let total_steps = if is_splitted_video {
            4
        } else if download.is_audio() {
            3
        } else {
            2
        };
        lib::update_steps(&downl_db.pool.clone(),&downl_db.qid, 2, total_steps);

        //download video, which is video/audio(m4a)
        try!(download.download_file(&file_path, false));

        if is_splitted_video { // download audio file & convert together
            lib::update_steps(&downl_db.pool.clone(),&downl_db.qid, 3, total_steps);

            let audio_path = format_audio_path(&downl_db.qid, folder.clone());
            file_db.push(audio_path.clone());

            println!("Downloading audio.. {}", audio_path);
            try!(download.download_file(&audio_path, true));

            lib::update_steps(&downl_db.pool.clone(),&downl_db.qid, 4, total_steps);

            match converter.merge_files(&downl_db.qid,&file_path, &audio_path,&save_path) {
                Err(e) => {println!("merge error: {:?}",e); return Err(e);},
                Ok(()) => {},
            }

            try!(remove_file(&audio_path));
            file_db.pop();

        }else{ // we're already done, only need to copy
            if !download.is_audio() {
                try!(lib::move_file(&file_path, &save_path));
            }
        }
    }
    if download.is_audio(){ // if audio-> convert m4a to mp3, which converts directly to downl. dir
        try!(converter.extract_audio(&downl_db.qid, &file_path, &save_path));
        try!(remove_file(&file_path));
    }else{
        if is_splitted_video {
            try!(remove_file(&file_path));
        }else{
            try!(lib::move_file(&file_path, &save_path));
        }
    }
    file_db.pop();

    if !is_zipped { // add file to list, except it's for zip-compression later (=folder set)
        lib::add_file_entry(&downl_db.pool.clone(), &downl_db.qid,&name_http_valid, &name);
    }

    //TODO: download file, convert if audio ?!
    Ok(true)
}

///Format save location for file, zip dependent
fn format_file_path(qid: &i64, folder: Option<String>) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}", &CONFIG.general.save_dir, v, qid),
        None => format!("{}/{}", &CONFIG.general.save_dir, qid),
    }
}

///Formats the audio path, based on the qid & optional folders
fn format_audio_path(qid: &i64, folder: Option<String>) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}a", &CONFIG.general.save_dir, v, qid),
        None => format!("{}/{}a", &CONFIG.general.save_dir, qid),
    }
}

///Format save path, dependent on zip option.
fn format_save_path<'a>(folder: Option<String>, name: &str, download: &'a Downloader, qid: &i64) -> String {
    match folder {
        Some(v) => format!("{}/{}/{}", &CONFIG.general.save_dir, v, lib::format_file_name(name,download,qid)),
        None => format!("{}/{}", &CONFIG.general.download_dir, lib::format_file_name(name,download,qid)),
    }
}

///Request an entry from the DB to handle
fn request_entry(pool: & pool::MyPool) -> Option<DownloadDB> {
    let mut conn = try_reoption!(pool.get_conn());
    let mut stmt = try_reoption!(conn.prepare("SELECT queries.qid,url,type,quality FROM querydetails \
                    INNER JOIN queries \
                    ON querydetails.qid = queries.qid \
                    WHERE querydetails.code = 0 \
                    ORDER BY queries.created \
                    LIMIT 1"));
    let mut result = try_reoption!(stmt.execute(&[]));
    let result = try_reoption!(try_option!(result.next())); // result.next().'Some'->value.'unwrap'

    println!("Result: {:?}", result[0]);
    println!("result str: {}", result[1].into_str());
    let download_db = DownloadDB { url: from_value::<String>(&result[1]),
                                    quality: from_value::<i16>(&result[3]),
                                    qid: from_value::<i64>(&result[0]),
                                    codecs: &CONFIG.codecs,
                                    extensions: &CONFIG.extensions,
                                    folder: CONFIG.general.save_dir.clone(),
                                    pool: pool.clone(),
                                    playlist: false, //TEMP
                                    compress: false };
    Some(download_db)
}

///Set dbms connection settings
fn mysql_options() -> MyOpts {
    MyOpts {
        tcp_addr: Some(CONFIG.db.ip.clone()),
        tcp_port: CONFIG.db.port,
        user: Some(CONFIG.db.user.clone()),
        pass: Some(CONFIG.db.password.clone()),
        db_name: Some(CONFIG.db.db.clone()),
        ..Default::default() // set others to default
    }
}<|endoftext|>


<commit_before>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::Result::{Ok, Err};
    use slice::SliceExt;
    use mem;
    use os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0u)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(v[mut read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as uint;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8, ..4] = [0u8, ..4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8, ..8] = [0u8, ..8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, INIT_ATOMIC_BOOL, Relaxed};

        static GETRANDOM_CHECKED: AtomicBool = INIT_ATOMIC_BOOL;
        static GETRANDOM_AVAILABLE: AtomicBool = INIT_ATOMIC_BOOL;

        if !GETRANDOM_CHECKED.load(Relaxed) {
            let mut buf: [u8, ..0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = errno() as libc::c_int;
                err != libc::ENOSYS
            } else {
                true
            };
            GETRANDOM_AVAILABLE.store(available, Relaxed);
            GETRANDOM_CHECKED.store(true, Relaxed);
            available
        } else {
            GETRANDOM_AVAILABLE.load(Relaxed)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use io::{IoResult};
    use kinds::marker;
    use mem;
    use os;
    use rand::Rng;
    use result::Result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::SliceExt;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    #[allow(missing_copy_implementations)]
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside of this module
        _dummy: (),
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    static kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::Result::{Ok, Err};
    use self::libc::{DWORD, BYTE, LPCSTR, BOOL};
    use self::libc::types::os::arch::extra::{LONG_PTR};
    use slice::SliceExt;

    type HCRYPTPROV = LONG_PTR;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[allow(non_snake_case)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0u, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0u, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(&mut v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<commit_msg>iOS: fallout of `marker::NoCopy` removal<commit_after>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::Result::{Ok, Err};
    use slice::SliceExt;
    use mem;
    use os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0u)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(v[mut read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as uint;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8, ..4] = [0u8, ..4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8, ..8] = [0u8, ..8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, INIT_ATOMIC_BOOL, Relaxed};

        static GETRANDOM_CHECKED: AtomicBool = INIT_ATOMIC_BOOL;
        static GETRANDOM_AVAILABLE: AtomicBool = INIT_ATOMIC_BOOL;

        if !GETRANDOM_CHECKED.load(Relaxed) {
            let mut buf: [u8, ..0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = errno() as libc::c_int;
                err != libc::ENOSYS
            } else {
                true
            };
            GETRANDOM_AVAILABLE.store(available, Relaxed);
            GETRANDOM_CHECKED.store(true, Relaxed);
            available
        } else {
            GETRANDOM_AVAILABLE.load(Relaxed)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use io::{IoResult};
    use mem;
    use os;
    use rand::Rng;
    use result::Result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::SliceExt;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    #[allow(missing_copy_implementations)]
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside of this module
        _dummy: (),
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    static kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::Result::{Ok, Err};
    use self::libc::{DWORD, BYTE, LPCSTR, BOOL};
    use self::libc::types::os::arch::extra::{LONG_PTR};
    use slice::SliceExt;

    type HCRYPTPROV = LONG_PTR;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[allow(non_snake_case)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0u, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0u, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(&mut v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<|endoftext|>


<commit_before>#![allow(non_upper_case_globals)]

use std::ffi::OsStr;
use std::path::PathBuf;

const build_props: &'static str = "project/build.properties";

fn home() -> PathBuf {
    std::env::home_dir().unwrap()
}

fn sbt_launch_dir() -> PathBuf {
    let mut p = PathBuf::from(home());
    p.push(".sbt/launchers");
    p
}

fn build_props_sbt() -> String {
    if let Ok(f) = std::fs::File::open(build_props) {
        let f = std::io::BufReader::new(f);

        use std::io::BufRead;
        for line in f.lines() {
            let line = line.unwrap();
            if line.starts_with("sbt.version") {
                let line = line.replace("=", " ");
                let line = line.replace("\r", " ");
                return line.split_whitespace().nth(1).unwrap().to_owned();
            }
        }
    }
    "".to_owned()
}

fn jar_file(sbt_version: &str) -> PathBuf {
    let mut p = PathBuf::from(sbt_launch_dir());
    p.push(sbt_version);
    p.push("sbt-launch.jar");
    p
}

fn exec_runner<S: AsRef<OsStr>>(args: &[S]) {
    use std::os::unix::process::CommandExt;
    let err = std::process::Command::new(&args[0]).args(&args[1..]).exec();
    println!("error: {}", err);
    if let Some(err) = err.raw_os_error() {
        std::process::exit(err);
    }
    std::process::exit(-1)
}

fn main() {
    let sbt_version = build_props_sbt();
    let sbt_jar = jar_file(&sbt_version);

    let java_cmd = "java";
    let extra_jvm_opts = ["-Xms512m", "-Xmx1536m", "-Xss2m"];
    let java_args: [&str; 0] = [];
    let sbt_commands: [&str; 0] = [];
    let residual_args = ["shell"];

    let mut exec_args: Vec<&OsStr> = Vec::new();
    exec_args.push(java_cmd.as_ref());
    exec_args.extend_from_slice(&extra_jvm_opts.iter().map(|x| x.as_ref()).collect::<Vec<_>>());
    exec_args.extend_from_slice(&java_args.iter().map(|x| x.as_ref()).collect::<Vec<_>>());
    exec_args.extend_from_slice(&["-jar".as_ref(), sbt_jar.as_ref()]);
    exec_args.extend_from_slice(&sbt_commands.iter().map(|x| x.as_ref()).collect::<Vec<_>>());
    exec_args.extend_from_slice(&residual_args.iter().map(|x| x.as_ref()).collect::<Vec<_>>());
    let exec_args = exec_args;

    exec_runner(&exec_args)
}
<commit_msg>Make build_props a PathBuf<commit_after>#![allow(non_upper_case_globals)]

use std::ffi::OsStr;
use std::path::PathBuf;

fn home() -> PathBuf {
    std::env::home_dir().unwrap()
}

fn sbt_launch_dir() -> PathBuf {
    let mut p = PathBuf::from(home());
    p.push(".sbt/launchers");
    p
}

fn build_props() -> PathBuf {
    PathBuf::from("project/build.properties")
}

fn build_props_sbt() -> String {
    if let Ok(f) = std::fs::File::open(build_props()) {
        let f = std::io::BufReader::new(f);

        use std::io::BufRead;
        for line in f.lines() {
            let line = line.unwrap();
            if line.starts_with("sbt.version") {
                let line = line.replace("=", " ");
                let line = line.replace("\r", " ");
                return line.split_whitespace().nth(1).unwrap().to_owned();
            }
        }
    }
    "".to_owned()
}

fn jar_file(sbt_version: &str) -> PathBuf {
    let mut p = PathBuf::from(sbt_launch_dir());
    p.push(sbt_version);
    p.push("sbt-launch.jar");
    p
}

fn exec_runner<S: AsRef<OsStr>>(args: &[S]) {
    use std::os::unix::process::CommandExt;
    let err = std::process::Command::new(&args[0]).args(&args[1..]).exec();
    println!("error: {}", err);
    if let Some(err) = err.raw_os_error() {
        std::process::exit(err);
    }
    std::process::exit(-1)
}

fn main() {
    let sbt_version = build_props_sbt();
    let sbt_jar = jar_file(&sbt_version);

    let java_cmd = "java";
    let extra_jvm_opts = ["-Xms512m", "-Xmx1536m", "-Xss2m"];
    let java_args: [&str; 0] = [];
    let sbt_commands: [&str; 0] = [];
    let residual_args = ["shell"];

    let mut exec_args: Vec<&OsStr> = Vec::new();
    exec_args.push(java_cmd.as_ref());
    exec_args.extend_from_slice(&extra_jvm_opts.iter().map(|x| x.as_ref()).collect::<Vec<_>>());
    exec_args.extend_from_slice(&java_args.iter().map(|x| x.as_ref()).collect::<Vec<_>>());
    exec_args.extend_from_slice(&["-jar".as_ref(), sbt_jar.as_ref()]);
    exec_args.extend_from_slice(&sbt_commands.iter().map(|x| x.as_ref()).collect::<Vec<_>>());
    exec_args.extend_from_slice(&residual_args.iter().map(|x| x.as_ref()).collect::<Vec<_>>());
    let exec_args = exec_args;

    exec_runner(&exec_args)
}
<|endoftext|>