<commit_before>#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper;
extern crate hubcaps;
extern crate afterparty;
extern crate serde_json;
extern crate serde;
extern crate rustc_serialize;
extern crate toml;

use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::env;
use std::io::Read;

use hyper::{Client, Server, Url};
use hyper::header::{Headers, Authorization, Bearer, UserAgent};
use hubcaps::{Credentials, Github};
use afterparty::*;
use serde_json::Value;
use regex::Regex;
use toml::decode_str;

#[derive(RustcEncodable, RustcDecodable)]
struct RepoData {
    pub qa_user: String,
    pub qa_flags: Vec<String>,
}

#[derive(RustcEncodable, RustcDecodable)]
struct ConfigData {
    pub api_root: String,
    pub repos: HashMap<String, RepoData>,
}

fn read_config_file(file_name: &str) -> ConfigData {
    let mut config_file = File::open(file_name).expect("Couldn't open config file");
    let mut str = String::new();
    config_file.read_to_string(&mut str).expect("Couldn't read config file");
    let toml: Option<ConfigData> = decode_str(&str);
    match toml {
        Some(cfg) => cfg,
        _ => panic!("Couldn't decode config file"),
    }
}

#[test]
fn it_reads_toml() {
    let config_data = read_config_file("test-data/test.toml");
    assert_eq!(config_data.api_root, "github.ibm.com/api/v3");
    let steve = config_data.repos.get("lkgrele/steve").unwrap();
    assert_eq!(steve.qa_user, "leif");
    assert_eq!(steve.qa_flags, ["qa", "test"]);
    let raul = config_data.repos.get("raul/robot").unwrap();
    assert_eq!(raul.qa_user, "raul");
    assert_eq!(raul.qa_flags, ["qa"]);
}

fn main() {
    let mut hub = Hub::new();
    hub.handle("pull_request", |delivery: &Delivery| {
        match delivery.payload {
            Event::PullRequest { number, ref pull_request, ref repository, .. } => {
                if pull_request.merged {
                    handle_pr(number, &pull_request.commits_url, &repository.full_name)
                }
            }
            _ => (),
        }
    });

    Server::http("0.0.0.0").unwrap().handle(hub).unwrap();
}

fn handle_pr(number: u64, commits_url: &str, repository: &str) {
    let config_data = read_config_file(".steve");
    let repo_config = config_data.repos.get(&repository.to_owned());
    let auth_token = env::var_os("STEVE_GITHUB_TOKEN")
        .expect("Missing github token!")
        .into_string()
        .unwrap();

    match repo_config {
        Some(config) => {
            update_issues(number,
                          &commits_url,
                          &auth_token,
                          &config,
                          &config_data.api_root,
                          &repository)
        }
        None => (), //just don't do anything if we don't have a config for this repo
    }
}

fn update_issues(number: u64,
                 commits_url: &str,
                 auth_token: &str,
                 config: &RepoData,
                 api_root: &str,
                 repository: &str) {
    let client = Client::new();
    let mut headers = Headers::new();
    headers.set(Authorization(Bearer { token: auth_token.to_owned() }));
    headers.set(UserAgent("steve".to_owned()));
    let commits = client.get(commits_url).headers(headers).send().unwrap();
    assert_eq!(commits.status, hyper::Ok); //replace with actual error handling

    match serde_json::from_reader(commits) {
        Ok(json) => {
            match parse_commit_data(&json) {
                Ok(messages) => {
                    let mut issues = HashSet::new();
                    for message in messages {
                        get_issues(&mut issues, &message)
                    }

                    for issue in issues {
                        update_issue(&client, auth_token, issue, config, &repository, api_root);
                    }
                }
                Err(err) => println!("{:?}", err), //should log error
            }
        }
        Err(err) => println!("{:?}", err), //should log error
    }
}

#[test]
fn it_parses_commit_data() {
    let file = File::open("test-data/test-messages.json").expect("Couldn't open message file");
    let json = serde_json::from_reader(file);
    assert!(json.is_ok());
    let json = json.unwrap();
    let messages = parse_commit_data(&json);
    assert!(messages.is_ok());
    let messages = messages.unwrap();
    assert_eq!(messages[0], "message1");
    assert_eq!(messages[1], "swordfish");
}

fn parse_commit_data<'a>(json: &'a Value) -> Result<Vec<&'a str>, &str> {
    let arr = try!(json.as_array().ok_or("expected array as top level object"));
    arr.iter()
        .map(|obj| {
            obj.pointer("/commit/message")
                .ok_or("missing message field")
                .and_then(|m| m.as_string().ok_or("message field has wrong type"))
        })
        .collect()
}

fn update_issue(client: &Client,
                auth_token: &str,
                issue_num: u64,
                config: &RepoData,
                repository: &str,
                api_root: &str) {
    let github = Github::host(api_root,
                              "steve",
                              client,
                              Credentials::Token(auth_token.to_owned()));
    let (owner, repo_name) = match repository.split("/").collect::<Vec<_>>() {
        strs => (strs[0], strs[1]),
    };
    let _ = github.repo(repo_name, owner)
        .issue(issue_num)
        .labels()
        .add(config.qa_flags.iter().map(String::as_str).collect());
    let mut headers = Headers::new();
    headers.set(Authorization(Bearer { token: auth_token.to_owned() }));
    headers.set(UserAgent("steve".to_owned()));
    let update_json = format!("{{\"assignee\":\"{}\"}}", config.qa_user);
    let url = Url::parse(&format!("{}/repos/{}/issues/{}", api_root, repository, issue_num))
        .expect("bad url");
    let res = client.patch(url)
        .body(&update_json)
        .headers(headers.to_owned())
        .send()
        .unwrap();
    assert_eq!(res.status, hyper::Ok);
}

#[test]
fn it_finds_issue_numbers() {
    let message = r"This is a commit message
        needs qa: #72 should get caught
        needs QA #333 should also get caught
        but #44 won't, and needs QA: #33 #45
        only gets #33. Duplicates are discarded:
        needs qa #72";
    let mut expected: HashSet<u64> = HashSet::new();
    expected.insert(72);
    expected.insert(333);
    expected.insert(33);
    let mut numbers: HashSet<u64> = HashSet::new();
    get_issues(&mut numbers, message);
    assert_eq!(numbers, expected);
}

fn get_issues(numbers: &mut HashSet<u64>, message: &str) {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"(?i)needs\sqa:?\s#(\d+)").unwrap();
    }
    for num in RE.captures_iter(message).filter_map(|cap| cap.at(1)) {
        numbers.insert(num.parse::<u64>().unwrap());
    }
}
<commit_msg>Remove unused "number" variable.<commit_after>#[macro_use]
extern crate lazy_static;
extern crate regex;
extern crate hyper;
extern crate hubcaps;
extern crate afterparty;
extern crate serde_json;
extern crate serde;
extern crate rustc_serialize;
extern crate toml;

use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::env;
use std::io::Read;

use hyper::{Client, Server, Url};
use hyper::header::{Headers, Authorization, Bearer, UserAgent};
use hubcaps::{Credentials, Github};
use afterparty::*;
use serde_json::Value;
use regex::Regex;
use toml::decode_str;

#[derive(RustcEncodable, RustcDecodable)]
struct RepoData {
    pub qa_user: String,
    pub qa_flags: Vec<String>,
}

#[derive(RustcEncodable, RustcDecodable)]
struct ConfigData {
    pub api_root: String,
    pub repos: HashMap<String, RepoData>,
}

fn read_config_file(file_name: &str) -> ConfigData {
    let mut config_file = File::open(file_name).expect("Couldn't open config file");
    let mut str = String::new();
    config_file.read_to_string(&mut str).expect("Couldn't read config file");
    let toml: Option<ConfigData> = decode_str(&str);
    match toml {
        Some(cfg) => cfg,
        _ => panic!("Couldn't decode config file"),
    }
}

#[test]
fn it_reads_toml() {
    let config_data = read_config_file("test-data/test.toml");
    assert_eq!(config_data.api_root, "github.ibm.com/api/v3");
    let steve = config_data.repos.get("lkgrele/steve").unwrap();
    assert_eq!(steve.qa_user, "leif");
    assert_eq!(steve.qa_flags, ["qa", "test"]);
    let raul = config_data.repos.get("raul/robot").unwrap();
    assert_eq!(raul.qa_user, "raul");
    assert_eq!(raul.qa_flags, ["qa"]);
}

fn main() {
    let mut hub = Hub::new();
    hub.handle("pull_request", |delivery: &Delivery| {
        match delivery.payload {
            Event::PullRequest { ref pull_request, ref repository, .. } => {
                if pull_request.merged {
                    handle_pr(&pull_request.commits_url, &repository.full_name)
                }
            }
            _ => (),
        }
    });

    Server::http("0.0.0.0").unwrap().handle(hub).unwrap();
}

fn handle_pr(commits_url: &str, repository: &str) {
    let config_data = read_config_file(".steve");
    let repo_config = config_data.repos.get(&repository.to_owned());
    let auth_token = env::var_os("STEVE_GITHUB_TOKEN")
        .expect("Missing github token!")
        .into_string()
        .unwrap();

    match repo_config {
        Some(config) => {
            update_issues(&commits_url,
                          &auth_token,
                          &config,
                          &config_data.api_root,
                          &repository)
        }
        None => (), //just don't do anything if we don't have a config for this repo
    }
}

fn update_issues(commits_url: &str,
                 auth_token: &str,
                 config: &RepoData,
                 api_root: &str,
                 repository: &str) {
    let client = Client::new();
    let mut headers = Headers::new();
    headers.set(Authorization(Bearer { token: auth_token.to_owned() }));
    headers.set(UserAgent("steve".to_owned()));
    let commits = client.get(commits_url).headers(headers).send().unwrap();
    assert_eq!(commits.status, hyper::Ok); //replace with actual error handling

    match serde_json::from_reader(commits) {
        Ok(json) => {
            match parse_commit_data(&json) {
                Ok(messages) => {
                    let mut issues = HashSet::new();
                    for message in messages {
                        get_issues(&mut issues, &message)
                    }

                    for issue in issues {
                        update_issue(&client, auth_token, issue, config, &repository, api_root);
                    }
                }
                Err(err) => println!("{:?}", err), //should log error
            }
        }
        Err(err) => println!("{:?}", err), //should log error
    }
}

#[test]
fn it_parses_commit_data() {
    let file = File::open("test-data/test-messages.json").expect("Couldn't open message file");
    let json = serde_json::from_reader(file);
    assert!(json.is_ok());
    let json = json.unwrap();
    let messages = parse_commit_data(&json);
    assert!(messages.is_ok());
    let messages = messages.unwrap();
    assert_eq!(messages[0], "message1");
    assert_eq!(messages[1], "swordfish");
}

fn parse_commit_data<'a>(json: &'a Value) -> Result<Vec<&'a str>, &str> {
    let arr = try!(json.as_array().ok_or("expected array as top level object"));
    arr.iter()
        .map(|obj| {
            obj.pointer("/commit/message")
                .ok_or("missing message field")
                .and_then(|m| m.as_string().ok_or("message field has wrong type"))
        })
        .collect()
}

fn update_issue(client: &Client,
                auth_token: &str,
                issue_num: u64,
                config: &RepoData,
                repository: &str,
                api_root: &str) {
    let github = Github::host(api_root,
                              "steve",
                              client,
                              Credentials::Token(auth_token.to_owned()));
    let (owner, repo_name) = match repository.split("/").collect::<Vec<_>>() {
        strs => (strs[0], strs[1]),
    };
    let _ = github.repo(repo_name, owner)
        .issue(issue_num)
        .labels()
        .add(config.qa_flags.iter().map(String::as_str).collect());
    let mut headers = Headers::new();
    headers.set(Authorization(Bearer { token: auth_token.to_owned() }));
    headers.set(UserAgent("steve".to_owned()));
    let update_json = format!("{{\"assignee\":\"{}\"}}", config.qa_user);
    let url = Url::parse(&format!("{}/repos/{}/issues/{}", api_root, repository, issue_num))
        .expect("bad url");
    let res = client.patch(url)
        .body(&update_json)
        .headers(headers.to_owned())
        .send()
        .unwrap();
    assert_eq!(res.status, hyper::Ok);
}

#[test]
fn it_finds_issue_numbers() {
    let message = r"This is a commit message
        needs qa: #72 should get caught
        needs QA #333 should also get caught
        but #44 won't, and needs QA: #33 #45
        only gets #33. Duplicates are discarded:
        needs qa #72";
    let mut expected: HashSet<u64> = HashSet::new();
    expected.insert(72);
    expected.insert(333);
    expected.insert(33);
    let mut numbers: HashSet<u64> = HashSet::new();
    get_issues(&mut numbers, message);
    assert_eq!(numbers, expected);
}

fn get_issues(numbers: &mut HashSet<u64>, message: &str) {
    lazy_static! {
        static ref RE: Regex = Regex::new(r"(?i)needs\sqa:?\s#(\d+)").unwrap();
    }
    for num in RE.captures_iter(message).filter_map(|cap| cap.at(1)) {
        numbers.insert(num.parse::<u64>().unwrap());
    }
}
<|endoftext|>


<commit_before>fn main() {

    let mut a = Box::new("abc");

    println!("main before abc run a is: {}", a);

    let b = abc(&mut a);

    println!("main after abc run a is {}", a);

    println!("main after abc run b is {}", b);
}

fn abc(a: &mut Box<&str>) -> &'static str {

    let temp: &str = get();

    *a = Box::new("abc changed");

    print_temp(temp);

    temp
}

fn get() -> &'static str {

    "xyz"
}

fn print_temp(a: &str) {

    println!("temp is: {}", a);
}<commit_msg>test: 提交修改<commit_after>fn main() {

    let mut a = "abc";

    println!("main before abc run a is: {}", a);

    let b = abc(&mut a);

    println!("main after abc run a is {}", a);

    println!("main after abc run b is {}", b);
}

fn abc(b: &mut &str) -> &'static str {

    let temp: &str = get();

    *b = "abc changed";

    print_temp(temp);

    temp
}

fn get() -> &'static str {

    "xyz"
}

fn print_temp(a: &str) {

    println!("temp is: {}", a);
}<|endoftext|>


<commit_before>extern crate combine;
extern crate nix; // As in *nix
extern crate rand;
extern crate signal;

mod commander;

// Reference: https://github.com/NixOS/nixpkgs/blob/b07051ce6c3ba3039c89b6755da279002b0c3ace/nixos/modules/virtualisation/nixos-container.pl

use nix::sys::signal::{kill, SIGTERM, SIGINT};
use std::io;
use std::io::Write;
use std::process::Command;

enum Error {
    Usage(String),
    NonZero(String, i32),
    ControlError(String),
    IoError(io::Error),
    NixError(nix::Error),
    Interrupted(commander::Signal)
}

use Error::*;

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { IoError(e) }
}

impl From<nix::Error> for Error {
    fn from(e: nix::Error) -> Self { NixError(e) }
}

impl From<commander::NonZero> for Error {
    fn from(e: commander::NonZero) -> Self { match e { commander::NonZero::NonZero(s, o) => NonZero(s, o), commander::NonZero::Interrupted(n) => Interrupted(n) } }
}

type Result<T> = std::result::Result<T, Error>;

// TODO: why can't this be inferred away
fn run<T>(cmd: &mut Command) -> Result<T> where T: commander::CommandOut { commander::run(cmd) }

fn probably_unique_name(length: usize) -> String {
    use rand::distributions::{IndependentSample, Range};
    let between = Range::new(0, 36);
    let mut rng = rand::thread_rng();
    unsafe {
        std::str::from_utf8_unchecked(&(0..length).map(|_| between.ind_sample(&mut rng)).map(|v| if v < 10 { v + 48 } else { v + 87 }).collect::<Vec<u8>>()).to_string()
    }
}

// Result is of the form "10.233.x"
fn unused_ip_block() -> Result<String> {
    // Can bail early anywhere here because we've got the lock
    let mut used = std::collections::HashSet::new();
    for entry in try!(std::fs::read_dir("/etc/containers")) {
        let entry = try!(entry);
        if try!(entry.file_type()).is_file() {
            let mut str = String::new();
            use std::io::Read;
            try!(try!(std::fs::File::open(entry.path())).read_to_string(&mut str));
            for line in str.lines() {
                use combine::{string, many1, choice, digit, token, eof, Parser, ParserExt};
                let ip_component = || { string("10.233.").with(many1::<String, _>(digit())).skip(token('.')).skip(many1::<String, _>(digit())) };
                if let Ok((c, _)) = choice([string("HOST_ADDRESS="), string("LOCAL_ADDRESS=")]).with(ip_component()).skip(eof()).parse(line) {
                    used.insert(c);
                }
            }
        }
    }
    (0..255).map(|p| p.to_string()).find(|s| !used.contains(s)).map(|s| format!("10.233.{}", s)).ok_or(ControlError("Out of IP addresses".to_string()))
}

fn mkpath(mode : u32, path : &str) -> Result<()> {
    use std::os::unix::fs::DirBuilderExt;
    match std::fs::DirBuilder::new().mode(mode).create(path) {
        Err(e) => if let std::io::ErrorKind::AlreadyExists = e.kind() { Ok(()) } else { Err(IoError(e)) },
        Ok(()) => Ok(())
    }
}

fn system_init() -> Result<()> {
    try!(mkpath(0o0755, "/etc/containers"));
    try!(mkpath(0o0700, "/var/lib/containers"));
    Ok(())
}

fn profile_dir(name : &str) -> String {
    format!("/nix/var/nix/profiles/per-container/{}", name)
}

fn container_root(name : &str) -> String {
    format!("/var/lib/containers/{}", name)
}

fn conf_file(name : &str) -> String {
    format!("/etc/containers/{}.conf", name)
}

fn create(config_file : &str) -> Result<String> {
    fn init_container_unsynced(name : &str) -> Result<()> {
        let ip_block = try!(unused_ip_block());
        // Use create_new to avoid clobbering existing container
        let mut f = try!(std::fs::OpenOptions::new().create_new(true).write(true).open(conf_file(name)));
        write!(&mut f, "PRIVATE_NETWORK=1\nHOST_ADDRESS={0}.1\nLOCAL_ADDRESS={0}.2\nAUTO_START=0\n", ip_block).map_err(IoError)
    }

    fn populate_container(name : &str, config_file : &str) -> Result<()> {
        use std::os::unix::fs::DirBuilderExt;
        // The per-container directory is restricted to prevent users on
        // the host from messing with guest users who happen to have the
        // same uid.
        try!(mkpath(0o0700, "/nix/var/nix/profiles/per-container"));
        let profile_dir = profile_dir(name);
        let mut dir_builder = std::fs::DirBuilder::new();
        try!(dir_builder.mode(0o0755).create(&profile_dir));

        let container_root = container_root(name);
        dir_builder.recursive(true);
        try!(dir_builder.mode(0o0755).create(format!("{}/etc/nixos", container_root)));

        // Write configuration.nix into the container fs
        let configuration_nix = format!("{}/etc/nixos/configuration.nix", container_root);
        {
            let mut f = try!(std::fs::File::create(format!("{}/etc/nixos/configuration.nix", container_root)));
            try!(write!(&mut f, "
    {{ config, lib, pkgs, ... }}:
    with lib;

    {{
        boot.isContainer = true;
        networking.hostName = mkDefault \"{}\";
        networking.useDHCP = false;
        imports = [ {} ];
    }}", name, config_file));
            // TODO: guaranteed sync at end of block?
        }

        // Build the container config
        run(Command::new("nix-env")
            .arg("-p").arg(format!("{}/system", profile_dir))
            .arg("-I").arg(format!("nixos-config={}", configuration_nix))
            .arg("-f").arg("<nixpkgs/nixos>")
            .arg("--set")
            .arg("-A").arg("system")
            .arg("--show-trace"))
    }

    let name = format!("fr{}", probably_unique_name(9));

    // Acquire lock
    use std::os::unix::io::IntoRawFd;
    let lock_fd = try!(std::fs::OpenOptions::new().create(true).append(true).open("/run/lock/nixos-container")).into_raw_fd();
    try!(nix::fcntl::flock(lock_fd, nix::fcntl::FlockArg::LockExclusive));

    let res = init_container_unsynced(&name);

    // Drop lock
    unsafe {
        let _ = nix::libc::close(lock_fd);  // If closing fails, ¯\_(ツ)_/¯
    }

    match res.and_then(|()| populate_container(&name, config_file)) {
        Ok(()) => Ok(name),
        // Don't leave a half-baked container behind
        Err(e) => { let _ = destroy(&name); Err(e) }
    }
}

fn run_test(container_name : &str, args : &[String]) -> Result<i32> {
    try!(run(Command::new("systemctl").arg("start").arg(format!("container@{}", container_name))));
    let leader_pid = try!(run(Command::new("machinectl").arg("show").arg(container_name).arg("-p").arg("Leader")).and_then(|o : String| {
        use combine::{string, many1, digit, spaces, eof, Parser, ParserExt};
        string("Leader=").with(many1::<String, _>(digit())).skip(spaces()).skip(eof()).parse(o.as_str())
            .map(|(ds, _)| ds)
            .map_err(|_| ControlError(format!("Bad machinectl output {}", o)))
    }));
    let su_cmd = args.iter().map(|s| format!("'{}'", s.replace("'", "'\\''"))).collect::<Vec<String>>().join(" ");
    match run(Command::new("nsenter")
            .arg("-t").arg(leader_pid)
            .arg("-m").arg("-u").arg("-i").arg("-n").arg("-p")
            .arg("--").arg("/var/setuid-wrappers/su").arg("root").arg("-l").arg("-c").arg(format!("exec {}", su_cmd))) {
        Ok(()) => Ok(0),
        Err(NonZero(_, code)) => Ok(code),
        Err(e) => Err(e)
    }
}

// Remove a directory while recursively unmounting all mounted filesystems within
// that directory and unmounting/removing that directory afterwards as well.
//
// Specified path shouldn't be a mountpoint.
fn safe_remove_tree(path : &str) -> Result<()> {
    if !std::path::Path::new(path).is_dir() {
        return Ok(());
    }
    try!(run(Command::new("find").arg(path)
        .arg("-mindepth").arg("1")
        .arg("-xdev")
        .arg("(")
            .arg("-type").arg("d")
            .arg("-exec").arg("mountpoint").arg("-q").arg("{}").arg(";")
        .arg(")")
        .arg("-exec").arg("umount").arg("-fR").arg("{}").arg("+")));
    try!(run(Command::new("rm").arg("--one-file-system").arg("-rf").arg(path)));

    Ok(())
}

fn stop(container_name : &str) -> Result<()> {
    match run(Command::new("systemctl").arg("stop").arg(format!("container@{}", container_name))) {
        Ok(()) => Ok(()),
        Err(Interrupted(_)) => {
            return run(Command::new("systemctl").arg("kill").arg(format!("container@{}", container_name)))
        },
        Err(e) => Err(e)
    }
}

fn destroy(container_name : &str) -> Result<()> {
    try!(stop(container_name));
    try!(safe_remove_tree(&profile_dir(container_name)));
    try!(safe_remove_tree(&format!("/nix/var/nix/gcroots/per-container/{}", container_name)));
    try!(safe_remove_tree(&container_root(container_name)));
    std::fs::remove_file(conf_file(container_name)).or_else(|e| match e.kind() {
        std::io::ErrorKind::NotFound => Ok(()),
        _ => Err(IoError(e))
    })
}

fn go() -> Result<i32> {
    let mut args = std::env::args();
    args.next(); // Discard program name
    let config_file = try!(args.next().ok_or(Error::Usage("Missing config file argument".to_string())));
    let test_args : Vec<String> = args.collect();
    if test_args.len() == 0 {
        return Err(Usage("Missing test args".to_string()));
    }

    let config_file = {
        let p = try!(std::fs::canonicalize(config_file).map_err(|e| Usage(format!("For config file: {}", e))));
        try!(p.to_str().ok_or(Usage("Config file path is invalid UTF-8".to_string()))).to_string()
    };

    try!(system_init());
    // We suppress signals in this process and allow the commander module to deal with child processes
    let trap = signal::trap::Trap::trap(&[SIGTERM, SIGINT]);
    let container_name = try!(create(&config_file));
    let res = match trap.wait(std::time::Instant::now()) {
        None => run_test(&container_name, &test_args),
        Some(n) => Err(Interrupted(n))
    };
    try!(stop(&container_name));

    res
}

fn main() {
    std::process::exit(go().unwrap_or_else(|e| match e {
        NonZero(cmd, code) => { writeln!(io::stderr(), "Command {:?} failed with code {:?}", cmd, code).unwrap(); 2 },
        Usage(msg) => { writeln!(io::stderr(), "Usage error: {}", msg).unwrap(); 1 },
        IoError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        NixError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        ControlError(msg) => { writeln!(io::stderr(), "{}", msg).unwrap(); 2 },
        Interrupted(SIGINT) => { kill(nix::unistd::getpid(), SIGINT).unwrap(); loop {} }
        Interrupted(n) => 128 + n
    }));
}
<commit_msg>Support --no_destroy argument and destroy by default<commit_after>extern crate combine;
extern crate nix; // As in *nix
extern crate rand;
extern crate signal;

mod commander;

// Reference: https://github.com/NixOS/nixpkgs/blob/b07051ce6c3ba3039c89b6755da279002b0c3ace/nixos/modules/virtualisation/nixos-container.pl

use nix::sys::signal::{kill, SIGTERM, SIGINT};
use std::io;
use std::io::Write;
use std::process::Command;

enum Error {
    Usage(String),
    NonZero(String, i32),
    ControlError(String),
    IoError(io::Error),
    NixError(nix::Error),
    Interrupted(commander::Signal)
}

use Error::*;

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { IoError(e) }
}

impl From<nix::Error> for Error {
    fn from(e: nix::Error) -> Self { NixError(e) }
}

impl From<commander::NonZero> for Error {
    fn from(e: commander::NonZero) -> Self { match e { commander::NonZero::NonZero(s, o) => NonZero(s, o), commander::NonZero::Interrupted(n) => Interrupted(n) } }
}

type Result<T> = std::result::Result<T, Error>;

// TODO: why can't this be inferred away
fn run<T>(cmd: &mut Command) -> Result<T> where T: commander::CommandOut { commander::run(cmd) }

fn probably_unique_name(length: usize) -> String {
    use rand::distributions::{IndependentSample, Range};
    let between = Range::new(0, 36);
    let mut rng = rand::thread_rng();
    unsafe {
        std::str::from_utf8_unchecked(&(0..length).map(|_| between.ind_sample(&mut rng)).map(|v| if v < 10 { v + 48 } else { v + 87 }).collect::<Vec<u8>>()).to_string()
    }
}

// Result is of the form "10.233.x"
fn unused_ip_block() -> Result<String> {
    // Can bail early anywhere here because we've got the lock
    let mut used = std::collections::HashSet::new();
    for entry in try!(std::fs::read_dir("/etc/containers")) {
        let entry = try!(entry);
        if try!(entry.file_type()).is_file() {
            let mut str = String::new();
            use std::io::Read;
            try!(try!(std::fs::File::open(entry.path())).read_to_string(&mut str));
            for line in str.lines() {
                use combine::{string, many1, choice, digit, token, eof, Parser, ParserExt};
                let ip_component = || { string("10.233.").with(many1::<String, _>(digit())).skip(token('.')).skip(many1::<String, _>(digit())) };
                if let Ok((c, _)) = choice([string("HOST_ADDRESS="), string("LOCAL_ADDRESS=")]).with(ip_component()).skip(eof()).parse(line) {
                    used.insert(c);
                }
            }
        }
    }
    (0..255).map(|p| p.to_string()).find(|s| !used.contains(s)).map(|s| format!("10.233.{}", s)).ok_or(ControlError("Out of IP addresses".to_string()))
}

fn mkpath(mode : u32, path : &str) -> Result<()> {
    use std::os::unix::fs::DirBuilderExt;
    match std::fs::DirBuilder::new().mode(mode).create(path) {
        Err(e) => if let std::io::ErrorKind::AlreadyExists = e.kind() { Ok(()) } else { Err(IoError(e)) },
        Ok(()) => Ok(())
    }
}

fn system_init() -> Result<()> {
    try!(mkpath(0o0755, "/etc/containers"));
    try!(mkpath(0o0700, "/var/lib/containers"));
    Ok(())
}

fn profile_dir(name : &str) -> String {
    format!("/nix/var/nix/profiles/per-container/{}", name)
}

fn container_root(name : &str) -> String {
    format!("/var/lib/containers/{}", name)
}

fn conf_file(name : &str) -> String {
    format!("/etc/containers/{}.conf", name)
}

fn create(config_file : &str) -> Result<String> {
    fn init_container_unsynced(name : &str) -> Result<()> {
        let ip_block = try!(unused_ip_block());
        // Use create_new to avoid clobbering existing container
        let mut f = try!(std::fs::OpenOptions::new().create_new(true).write(true).open(conf_file(name)));
        write!(&mut f, "PRIVATE_NETWORK=1\nHOST_ADDRESS={0}.1\nLOCAL_ADDRESS={0}.2\nAUTO_START=0\n", ip_block).map_err(IoError)
    }

    fn populate_container(name : &str, config_file : &str) -> Result<()> {
        use std::os::unix::fs::DirBuilderExt;
        // The per-container directory is restricted to prevent users on
        // the host from messing with guest users who happen to have the
        // same uid.
        try!(mkpath(0o0700, "/nix/var/nix/profiles/per-container"));
        let profile_dir = profile_dir(name);
        let mut dir_builder = std::fs::DirBuilder::new();
        try!(dir_builder.mode(0o0755).create(&profile_dir));

        let container_root = container_root(name);
        dir_builder.recursive(true);
        try!(dir_builder.mode(0o0755).create(format!("{}/etc/nixos", container_root)));

        // Write configuration.nix into the container fs
        let configuration_nix = format!("{}/etc/nixos/configuration.nix", container_root);
        {
            let mut f = try!(std::fs::File::create(format!("{}/etc/nixos/configuration.nix", container_root)));
            try!(write!(&mut f, "
    {{ config, lib, pkgs, ... }}:
    with lib;

    {{
        boot.isContainer = true;
        networking.hostName = mkDefault \"{}\";
        networking.useDHCP = false;
        imports = [ {} ];
    }}", name, config_file));
            // TODO: guaranteed sync at end of block?
        }

        // Build the container config
        run(Command::new("nix-env")
            .arg("-p").arg(format!("{}/system", profile_dir))
            .arg("-I").arg(format!("nixos-config={}", configuration_nix))
            .arg("-f").arg("<nixpkgs/nixos>")
            .arg("--set")
            .arg("-A").arg("system")
            .arg("--show-trace"))
    }

    let name = format!("fr{}", probably_unique_name(9));

    // Acquire lock
    use std::os::unix::io::IntoRawFd;
    let lock_fd = try!(std::fs::OpenOptions::new().create(true).append(true).open("/run/lock/nixos-container")).into_raw_fd();
    try!(nix::fcntl::flock(lock_fd, nix::fcntl::FlockArg::LockExclusive));

    let res = init_container_unsynced(&name);

    // Drop lock
    unsafe {
        let _ = nix::libc::close(lock_fd);  // If closing fails, ¯\_(ツ)_/¯
    }

    match res.and_then(|()| populate_container(&name, config_file)) {
        Ok(()) => Ok(name),
        // Don't leave a half-baked container behind
        Err(e) => { let _ = destroy(&name); Err(e) }
    }
}

fn run_test(container_name : &str, args : &[String]) -> Result<i32> {
    try!(run(Command::new("systemctl").arg("start").arg(format!("container@{}", container_name))));
    let leader_pid = try!(run(Command::new("machinectl").arg("show").arg(container_name).arg("-p").arg("Leader")).and_then(|o : String| {
        use combine::{string, many1, digit, spaces, eof, Parser, ParserExt};
        string("Leader=").with(many1::<String, _>(digit())).skip(spaces()).skip(eof()).parse(o.as_str())
            .map(|(ds, _)| ds)
            .map_err(|_| ControlError(format!("Bad machinectl output {}", o)))
    }));
    let su_cmd = args.iter().map(|s| format!("'{}'", s.replace("'", "'\\''"))).collect::<Vec<String>>().join(" ");
    match run(Command::new("nsenter")
            .arg("-t").arg(leader_pid)
            .arg("-m").arg("-u").arg("-i").arg("-n").arg("-p")
            .arg("--").arg("/var/setuid-wrappers/su").arg("root").arg("-l").arg("-c").arg(format!("exec {}", su_cmd))) {
        Ok(()) => Ok(0),
        Err(NonZero(_, code)) => Ok(code),
        Err(e) => Err(e)
    }
}

// Remove a directory while recursively unmounting all mounted filesystems within
// that directory and unmounting/removing that directory afterwards as well.
//
// Specified path shouldn't be a mountpoint.
fn safe_remove_tree(path : &str) -> Result<()> {
    if !std::path::Path::new(path).is_dir() {
        return Ok(());
    }
    try!(run(Command::new("find").arg(path)
        .arg("-mindepth").arg("1")
        .arg("-xdev")
        .arg("(")
            .arg("-type").arg("d")
            .arg("-exec").arg("mountpoint").arg("-q").arg("{}").arg(";")
        .arg(")")
        .arg("-exec").arg("umount").arg("-fR").arg("{}").arg("+")));
    try!(run(Command::new("rm").arg("--one-file-system").arg("-rf").arg(path)));

    Ok(())
}

fn stop(container_name : &str) -> Result<()> {
    match run(Command::new("systemctl").arg("stop").arg(format!("container@{}", container_name))) {
        Ok(()) => Ok(()),
        Err(Interrupted(_)) => {
            return run(Command::new("systemctl").arg("kill").arg(format!("container@{}", container_name)))
        },
        Err(e) => Err(e)
    }
}

fn destroy(container_name : &str) -> Result<()> {
    try!(stop(container_name));
    try!(safe_remove_tree(&profile_dir(container_name)));
    try!(safe_remove_tree(&format!("/nix/var/nix/gcroots/per-container/{}", container_name)));
    try!(safe_remove_tree(&container_root(container_name)));
    std::fs::remove_file(conf_file(container_name)).or_else(|e| match e.kind() {
        std::io::ErrorKind::NotFound => Ok(()),
        _ => Err(IoError(e))
    })
}

fn go() -> Result<i32> {
    // Skip program name
    let mut args = std::env::args().skip(1).peekable();

    let no_destroy = args.peek().map_or(false, |arg| arg == "--no_destroy");
    if no_destroy { args.next(); }

    let config_file = try!(args.next().ok_or(Error::Usage("Missing config file argument".to_string())));

    let test_args : Vec<String> = args.collect();
    if test_args.len() == 0 {
        return Err(Usage("Missing test args".to_string()));
    }

    let config_file = {
        let p = try!(std::fs::canonicalize(config_file).map_err(|e| Usage(format!("For config file: {}", e))));
        try!(p.to_str().ok_or(Usage("Config file path is invalid UTF-8".to_string()))).to_string()
    };

    try!(system_init());
    // We suppress signals in this process and allow the commander module to deal with child processes
    let trap = signal::trap::Trap::trap(&[SIGTERM, SIGINT]);
    let container_name = try!(create(&config_file));
    let res = match trap.wait(std::time::Instant::now()) {
        None => run_test(&container_name, &test_args),
        Some(n) => Err(Interrupted(n))
    };

    if no_destroy {
        try!(stop(&container_name));
    } else {
        try!(destroy(&container_name));
    }

    res
}

fn main() {
    std::process::exit(go().unwrap_or_else(|e| match e {
        NonZero(cmd, code) => { writeln!(io::stderr(), "Command {:?} failed with code {:?}", cmd, code).unwrap(); 2 },
        Usage(msg) => { writeln!(io::stderr(), "Usage error: {}", msg).unwrap(); 1 },
        IoError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        NixError(err) => { writeln!(io::stderr(), "{}", err).unwrap(); 2 },
        ControlError(msg) => { writeln!(io::stderr(), "{}", msg).unwrap(); 2 },
        Interrupted(SIGINT) => { kill(nix::unistd::getpid(), SIGINT).unwrap(); loop {} }
        Interrupted(n) => 128 + n
    }));
}
<|endoftext|>


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

extern crate getopts;
extern crate char_iter;

mod trtable;

use std::clone::Clone;
use std::io;
use std::io::Write;
use std::{env, thread};
use std::sync::mpsc::{channel, Receiver, Sender};
use getopts::Options;
use trtable::TrTable;

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

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

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

fn all_combos(tx: &Sender<String>, chars: &Vec<char>, count: usize, built: String) {
  // base case
  if count == 0 {
    //drop(tx.send(&built));
    let stdout = io::stdout();
    writeln!(&mut stdout.lock(), "{}", built);
    //guard.flush();
    return;
  }

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

fn print_usage(program: &str, opts: Options) {
  let brief = format!("Usage: {} [options]", program);
  print!("{}", opts.usage(&brief));
}

fn main() {
  let args: Vec<String> = env::args().collect();
  let program = args[0].clone();

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

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

  if matches.opt_present("h") {
    // print help and exit
    print_usage(&program, opts);
    return;
  }

  let table_file = matches.opt_str("t");
  let num_threads = match matches.opt_str("j") {
    Some(n_jobs) => n_jobs.parse::<usize>().unwrap(),
    None => 1
  };

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

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

    // threads that are currently spawned
    let mut receivers: Vec<Receiver<String>> = Vec::new();
    let mut received = 0;

    for n in counts {
      for letter in &all_chars {
        // get whether we have any thread slots open; if not, then
        // block until a thread finishes.
        while receivers.len() == num_threads {
          let mut offset = 0;
          for index in 0 .. receivers.len() {
            let remove = match receivers[index - offset].try_recv() {
              Ok(s) => if s == "" {
                true
              } else {
                println!("{}", s);
                false
              },
              Err(_) => false,
            };

            if remove {
              receivers.remove(index - offset);
            }
          }
          thread::sleep_ms(1);
        }

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

          new_built.push(l);
          all_combos(&tx, &all_chars_dup, n - 1, new_built);
          tx.send(String::new());
        });
        receivers.push(rx);
      }
    }
  }
}
<commit_msg>Cleaned up the code some, removed some warnings, etc<commit_after>/*
 * Copyright (C) 2015 Alek Ratzloff
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

extern crate getopts;
extern crate char_iter;

mod trtable;

use std::clone::Clone;
use std::io;
use std::io::Write;
use std::{env, thread};
use std::sync::mpsc::{channel, Receiver};
use getopts::Options;
use trtable::TrTable;

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

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

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

fn all_combos(chars: &Vec<char>, count: usize, built: String) {
  // base case
  if count == 0 {
    //drop(tx.send(&built));
    let stdout = io::stdout();
    drop(writeln!(&mut stdout.lock(), "{}", built));
    //guard.flush();
    return;
  }

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

fn print_usage(program: &str, opts: Options) {
  let brief = format!("Usage: {} [options]", program);
  print!("{}", opts.usage(&brief));
}

fn main() {
  let args: Vec<String> = env::args().collect();
  let program = args[0].clone();

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

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

  if matches.opt_present("h") {
    // print help and exit
    print_usage(&program, opts);
    return;
  }

  let table_file = matches.opt_str("t");
  let num_threads = match matches.opt_str("j") {
    Some(n_jobs) => n_jobs.parse::<usize>().unwrap(),
    None => 1
  };

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

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

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

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

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

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

          new_built.push(l);
          all_combos(&all_chars_dup, n - 1, new_built);
          drop(tx.send(0));
        });
        receivers.push(rx);
      }
    }
  }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate error_chain;
#[macro_use]
extern crate stderr;
extern crate chan;
extern crate clap;
extern crate filetime;
extern crate glob;

#[macro_use] mod errors;
mod compression;
mod lists;
mod structs;

use clap::{App, Arg};
use errors::*;
use lists::*;
use structs::*;
use std::path::{Path, PathBuf};
use std::sync::Arc;

quick_main!(run);

fn run() -> Result<()> {
    let matches = App::new("static-compress")
        .version("0.1")
        .about("Create statically-compresed copies of matching files")
        .author("NeoSmart Technologies")
        .arg(Arg::with_name("compressor")
            .short("c")
            .long("compressor")
            .value_name("[brotli|gzip|zopfli]")
            .help("The compressor to use (default: gzip)")
            .takes_value(true))
        .arg(Arg::with_name("threads")
            .short("j")
            .long("threads")
            .value_name("COUNT")
            .help("The number of simultaneous compressions (default: 1)")
            .takes_value(true))
        .arg(Arg::with_name("filters")
            .value_name("FILTER")
            .multiple(true)
            .required(true))
        .arg(Arg::with_name("ext")
            .short("e")
            .value_name("EXT")
            .long("extension")
            .help("The extension to use for compressed files (default: gz or br)"))
        /*.arg(Arg::with_name("excludes")
            .short("x")
            .value_name("FILTER")
            .long("exclude")
            .multiple(true)
            .help("Exclude files matching this glob expression"))*/
        .get_matches();

    fn get_parameter<'a, T>(matches: &clap::ArgMatches, name: &str, default_value: T) -> Result<T>
        where T: std::str::FromStr
    {
        match matches.value_of(name) {
            Some(v) => {
                Ok(v.parse().map_err(|_| ErrorKind::InvalidParameterValue(name.to_owned()))?)
            }
            None => Ok(default_value),
        }
    }

    let compressor = get_parameter(&matches, "compressor", CompressionAlgorithm::GZip)?;
    let temp = Parameters {
        extension: matches.value_of("ext")
            .unwrap_or(compressor.extension())
            .trim_matches(|c: char| c.is_whitespace() || c.is_control() || c == '.')
            .to_owned(),
        compressor: compressor,
        include_filters: match matches.values_of("filters") {
            Some(values) => Ok(values.map(|s| s.to_owned()).collect()),
            None => Err(ErrorKind::InvalidUsage),
        }?,
        threads: get_parameter(&matches, "threads", 1)?,
    };

    /*let exclude_filters = match matches.values_of("exclude") {
        Some(values)=> values.map(|s| s.to_owned()).collect(),
        None => Vec::<String>::new(),
    };*/

    let parameters = Arc::<Parameters>::new(temp);
    let (send_queue, wait_group) = start_workers(&parameters);

    //convert filters to paths and deal out conversion jobs
    dispatch_jobs(send_queue, &parameters.include_filters/*, exclude_filters*/)?;

    //wait for all jobs to finish
    wait_group.wait();

    Ok(())
}

type ThreadParam = std::path::PathBuf;

fn start_workers<'a>(params: &Arc<Parameters>) -> (chan::Sender<ThreadParam>, chan::WaitGroup) {
    let (tx, rx) = chan::sync::<ThreadParam>(params.threads);
    let wg = chan::WaitGroup::new();

    for _ in 0..params.threads {
        let local_params = params.clone();
        let local_rx = rx.clone();
        let local_wg = wg.clone();
        wg.add(1);
        std::thread::spawn(move || {
            worker_thread(local_params, local_rx);
            local_wg.done();
        });
    }

    (tx, wg)
}

fn dispatch_jobs(send_queue: chan::Sender<ThreadParam>, filters: &Vec<String>/*, exclude_filters: Vec<String>*/) -> Result<()> {
    for filter in filters {
        for entry in glob::glob(filter).map_err(|_| ErrorKind::InvalidIncludeFilter)? {
            match entry {
                Ok(path) => {
                    if is_blacklisted(&path)? {
                        //this path has been excluded
                        continue;
                    }
                    //make sure this is a file, not a folder
                    match std::fs::metadata(&path) {
                        Ok(metadata) => {
                            if metadata.is_file() {
                                send_queue.send(path);
                            }
                            continue; //skip otherwise
                        }
                        Err(e) => errstln!("{}: {}", path.to_string_lossy(), e),
                    }
                }
                Err(e) => errstln!("{}", e),
            };
        }
    }
    Ok(())
}

fn worker_thread(params: Arc<Parameters>, rx: chan::Receiver<ThreadParam>) {
    loop {
        let src = match rx.recv() {
            Some(task) => task,
            None => return, //no more tasks
        };

        //in a nested function so we can handle errors centrally
        fn compress_single(src: &ThreadParam, params: &Parameters) -> Result<()> {
            let dst_path = format!("{}.{}",
                                   src.to_str().ok_or(ErrorKind::InvalidCharactersInPath)?,
                                   params.extension);
            let dst = Path::new(&dst_path);

            //again, in a scope for error handling
            || -> Result<()> {
                    println!("{}", src.to_string_lossy());
                    let src_metadata = std::fs::metadata(src)?;

                    //don't compress files that are already compressed that haven't changed
                    if let Ok(dst_metadata) = std::fs::metadata(dst) {
                        //the destination already exists
                        match src_metadata.modified()? == dst_metadata.modified()? {
                            true => return Ok(()), //no need to recompress
                            false => std::fs::remove_file(dst)?, //throw if we can't
                        };
                    }
                    params.compressor.compress(src.as_path(), dst)?;
                    let src_modified = filetime::FileTime::from_last_modification_time(&src_metadata);
                    filetime::set_file_times(dst, filetime::FileTime::zero(), src_modified).unwrap_or_default();

                    Ok(())
                }()
                .map_err(|e| {
                    //try deleting the invalid destination file, but don't care if we can't
                    std::fs::remove_file(dst).unwrap_or_default();
                    e //return the same error
                })
        }

        if let Err(e) = compress_single(&src, &params) {
            errstln!("Error compressing {}: {}", src.to_string_lossy(), e);
        }
    }
}

fn str_search(sorted: &[&str], search_term: &str, case_sensitive: bool) -> std::result::Result<usize, usize> {
    let term = match case_sensitive {
        true => search_term.to_owned(),
        false => search_term.to_lowercase(),
    };

    sorted.binary_search_by(|probe| probe.cmp(&&*term))
}

fn is_blacklisted(path: &PathBuf) -> Result<bool> {
    let r = match path.extension() {
        Some(x) => {
            let ext = x.to_str().ok_or(ErrorKind::InvalidCharactersInPath)?;
            str_search(COMP_EXTS, &ext, false).is_ok()
        },
        None => false,
    };

    return Ok(r);
}

<commit_msg>Support globbing all files via only **<commit_after>#[macro_use]
extern crate error_chain;
#[macro_use]
extern crate stderr;
extern crate chan;
extern crate clap;
extern crate filetime;
extern crate glob;

#[macro_use] mod errors;
mod compression;
mod lists;
mod structs;

use clap::{App, Arg};
use errors::*;
use lists::*;
use structs::*;
use std::path::{Path, PathBuf};
use std::sync::Arc;

quick_main!(run);

fn run() -> Result<()> {
    let matches = App::new("static-compress")
        .version("0.1")
        .about("Create statically-compresed copies of matching files")
        .author("NeoSmart Technologies")
        .arg(Arg::with_name("compressor")
            .short("c")
            .long("compressor")
            .value_name("[brotli|gzip|zopfli]")
            .help("The compressor to use (default: gzip)")
            .takes_value(true))
        .arg(Arg::with_name("threads")
            .short("j")
            .long("threads")
            .value_name("COUNT")
            .help("The number of simultaneous compressions (default: 1)")
            .takes_value(true))
        .arg(Arg::with_name("filters")
            .value_name("FILTER")
            .multiple(true)
            .required(true))
        .arg(Arg::with_name("ext")
            .short("e")
            .value_name("EXT")
            .long("extension")
            .help("The extension to use for compressed files (default: gz or br)"))
        /*.arg(Arg::with_name("excludes")
            .short("x")
            .value_name("FILTER")
            .long("exclude")
            .multiple(true)
            .help("Exclude files matching this glob expression"))*/
        .get_matches();

    fn get_parameter<'a, T>(matches: &clap::ArgMatches, name: &str, default_value: T) -> Result<T>
        where T: std::str::FromStr
    {
        match matches.value_of(name) {
            Some(v) => {
                Ok(v.parse().map_err(|_| ErrorKind::InvalidParameterValue(name.to_owned()))?)
            }
            None => Ok(default_value),
        }
    }

    let compressor = get_parameter(&matches, "compressor", CompressionAlgorithm::GZip)?;
    let temp = Parameters {
        extension: matches.value_of("ext")
            .unwrap_or(compressor.extension())
            .trim_matches(|c: char| c.is_whitespace() || c.is_control() || c == '.')
            .to_owned(),
        compressor: compressor,
        include_filters: match matches.values_of("filters") {
            Some(values) => Ok(values.map(|s| s.to_owned()).collect()),
            None => Err(ErrorKind::InvalidUsage),
        }?,
        threads: get_parameter(&matches, "threads", 1)?,
    };

    /*let exclude_filters = match matches.values_of("exclude") {
        Some(values)=> values.map(|s| s.to_owned()).collect(),
        None => Vec::<String>::new(),
    };*/

    let parameters = Arc::<Parameters>::new(temp);
    let (send_queue, wait_group) = start_workers(&parameters);

    //convert filters to paths and deal out conversion jobs
    dispatch_jobs(send_queue, &parameters.include_filters/*, exclude_filters*/)?;

    //wait for all jobs to finish
    wait_group.wait();

    Ok(())
}

type ThreadParam = std::path::PathBuf;

fn start_workers<'a>(params: &Arc<Parameters>) -> (chan::Sender<ThreadParam>, chan::WaitGroup) {
    let (tx, rx) = chan::sync::<ThreadParam>(params.threads);
    let wg = chan::WaitGroup::new();

    for _ in 0..params.threads {
        let local_params = params.clone();
        let local_rx = rx.clone();
        let local_wg = wg.clone();
        wg.add(1);
        std::thread::spawn(move || {
            worker_thread(local_params, local_rx);
            local_wg.done();
        });
    }

    (tx, wg)
}

fn dispatch_jobs(send_queue: chan::Sender<ThreadParam>, filters: &Vec<String>/*, exclude_filters: Vec<String>*/) -> Result<()> {
    for filter in filters {
        //by default, rs-glob treats "**" as a directive to match only directories
        //we can either rewrite "**" as "**/*" or recurse into directories below
        let new_filter = (&*filter).replace("**", "**/*");
        for entry in glob::glob(&new_filter).map_err(|_| ErrorKind::InvalidIncludeFilter)? {
            match entry {
                Ok(path) => {
                    if is_blacklisted(&path)? {
                        //this path has been excluded
                        continue;
                    }
                    //make sure this is a file, not a folder
                    match std::fs::metadata(&path) {
                        Ok(metadata) => {
                            if metadata.is_file() {
                                send_queue.send(path);
                            }
                            continue; //skip otherwise
                        }
                        Err(e) => errstln!("{}: {}", path.to_string_lossy(), e),
                    }
                }
                Err(e) => errstln!("{}", e),
            };
        }
    }
    Ok(())
}

fn worker_thread(params: Arc<Parameters>, rx: chan::Receiver<ThreadParam>) {
    loop {
        let src = match rx.recv() {
            Some(task) => task,
            None => return, //no more tasks
        };

        //in a nested function so we can handle errors centrally
        fn compress_single(src: &ThreadParam, params: &Parameters) -> Result<()> {
            let dst_path = format!("{}.{}",
                                   src.to_str().ok_or(ErrorKind::InvalidCharactersInPath)?,
                                   params.extension);
            let dst = Path::new(&dst_path);

            //again, in a scope for error handling
            || -> Result<()> {
                    println!("{}", src.to_string_lossy());
                    let src_metadata = std::fs::metadata(src)?;

                    //don't compress files that are already compressed that haven't changed
                    if let Ok(dst_metadata) = std::fs::metadata(dst) {
                        //the destination already exists
                        match src_metadata.modified()? == dst_metadata.modified()? {
                            true => return Ok(()), //no need to recompress
                            false => std::fs::remove_file(dst)?, //throw if we can't
                        };
                    }
                    params.compressor.compress(src.as_path(), dst)?;
                    let src_modified = filetime::FileTime::from_last_modification_time(&src_metadata);
                    filetime::set_file_times(dst, filetime::FileTime::zero(), src_modified).unwrap_or_default();

                    Ok(())
                }()
                .map_err(|e| {
                    //try deleting the invalid destination file, but don't care if we can't
                    std::fs::remove_file(dst).unwrap_or_default();
                    e //return the same error
                })
        }

        if let Err(e) = compress_single(&src, &params) {
            errstln!("Error compressing {}: {}", src.to_string_lossy(), e);
        }
    }
}

fn str_search(sorted: &[&str], search_term: &str, case_sensitive: bool) -> std::result::Result<usize, usize> {
    let term = match case_sensitive {
        true => search_term.to_owned(),
        false => search_term.to_lowercase(),
    };

    sorted.binary_search_by(|probe| probe.cmp(&&*term))
}

fn is_blacklisted(path: &PathBuf) -> Result<bool> {
    let r = match path.extension() {
        Some(x) => {
            let ext = x.to_str().ok_or(ErrorKind::InvalidCharactersInPath)?;
            str_search(COMP_EXTS, &ext, false).is_ok()
        },
        None => false,
    };

    return Ok(r);
}

<|endoftext|>


<commit_before>#![feature(core)]
#![feature(path)]

#![feature(plugin)]
#![plugin(regex_macros)]
extern crate regex;
extern crate git2;
extern crate core;
extern crate "rustc-serialize" as rustc_serialize;
extern crate docopt;

use git2::{Repository, Branch, BranchType, DiffLine,
    Commit, Diff, DiffFormat, DiffDelta, DiffHunk, Oid};
use docopt::Docopt;
use std::old_io as io;
use std::collections::HashMap;
use std::old_io::BufferedReader;
use std::old_io::File;
use std::str;
use regex::Regex;
use rustc_serialize::json::{self, ToJson, Json};
use std::fmt::Display;

// Write the Docopt usage string.
static USAGE: &'static str = "
Usage: transit <repo> [--json] [<old> <new>]

If no commits are given, transit will revwalk from latest to oldest.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_repo: String,
    arg_old: Option<String>,
    arg_new: Option<String>,
    flag_json: bool,
}

fn main() {
    // Parse the args above or die.
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());
    // Pull up the revwalk.
    let path = Path::new(args.arg_repo);
    let repo = Repository::discover(&path)
        .ok().expect("Unable to find repo.");
    // If we can destructures the two optional args into real things.
    if let (Some(old_string), Some(new_string)) = (args.arg_old, args.arg_new) {
        // Compare a specific commit pair.
        // Fist, get the commits. (Error checked)
        let old = Oid::from_str(&old_string[]).and_then(|oid| repo.find_commit(oid));
        let new = Oid::from_str(&new_string[]).and_then(|oid| repo.find_commit(oid));
        if old.is_ok() && new.is_ok() {
            let output = find_moves(&repo, &old.unwrap(), &new.unwrap()).unwrap();
            make_output(output);
        } else {
            panic!("Commit ids were not valid.");
        }
    } else {
        // Revwalk.
        let mut revwalk = repo.revwalk()
            .ok().expect("Unable to get revwalk.");
        // Setup some options.
        revwalk.simplify_first_parent(); // TODO: Maybe remove?
        let mut flags = git2::Sort::empty();
        flags.insert(git2::SORT_TIME);
        flags.insert(git2::SORT_TOPOLOGICAL);
        revwalk.set_sorting(flags);
        // Push HEAD to the revwalk.
        revwalk.push_head()
            .ok().expect("Unable to push HEAD.");
        // We sadly must collect here to use `.windows()`
        let history = revwalk.filter_map(|id| repo.find_commit(id).ok())
            .collect::<Vec<Commit>>();
        // Walk through each pair of commits.
        for pair in history.windows(2) {let (old, new) = (&pair[1], &pair[0]);
            let output = find_moves(&repo, old.clone(), new.clone()).unwrap();
            if args.flag_json {
                make_json(output);
            } else {
                make_output(output);
            }
        }
    }
}

fn make_output(output: Vec<Output>) {
    println!("make_output: output.len()={}", output.len());
    for i in range(0, output.len()) {
            println!("\told_commit={}",         output[i].old_commit);
            println!("\tnew_commit={}",         output[i].new_commit);
            println!("\torigin_line={}",        output[i].origin_line);
            println!("\tdestintation_line={}",  output[i].destination_line);
            println!("\tnum_lines={}",          output[i].num_lines);
            println!("\tnew_filename={}",       output[i].new_filename);
            println!("\told_filename={}",       output[i].old_filename);
    }
}

fn make_json(output: Vec<Output>) {
    println!("{}", json::encode(&output).unwrap());
}

#[derive(Debug, PartialEq, Eq)]
enum FoundState {
    Added, Deleted
}

fn dump_diffline(line: &DiffLine) {
    // 'origin' is wrapped in pipes to ease displaying space characters.
    print!("line: old={:?} new={:?} offset={} |origin|=|{}|\n      content={}",
             line.old_lineno(), line.new_lineno(), line.content_offset(),
             line.origin(), str::from_utf8(line.content()).unwrap());
}

fn dump_diffdelta(delta: &DiffDelta) {
    println!("delta: nfiles={} status={:?} old_file=(id={} path_bytes={:?} path={:?} tsize={}) new_file=(id={} path_bytes={:?} path={:?} tsize={})",
            delta.nfiles(), delta.status(),
            delta.old_file().id(), delta.old_file().path_bytes(), delta.old_file().path(), delta.old_file().size(),
            delta.new_file().id(), delta.new_file().path_bytes(), delta.new_file().path(), delta.new_file().size());
}

fn dump_diffhunk(hunk: &DiffHunk) {
    println!("hunk: old_start={} old_lines={} new_start={} new_lines={} header={}",
            hunk.old_start(), hunk.old_lines(),
            hunk.new_start(), hunk.new_lines(),
            str::from_utf8(hunk.header()).unwrap());
}

fn format_key(key: String) -> String {
    let remove_whitespace = regex!(r"\s{2,}"); // 2 or more whitespaces    // TODO Removes whitespace from a string.
    let trim = regex!(r"^[\s]+|[\s]+$");
    let result = remove_whitespace.replace_all(key.as_slice(), "");
    trim.replace_all(result.as_slice(), "")
}

#[derive(Debug)]
struct Found {
    filename: Path,
    key: String,
    state: FoundState,
    start_position: u32,
    line_count: u32,
}

fn find_additions_and_deletions(diff: Diff) -> Vec<Found> {

    #[derive(Debug)]
    enum State {
        Other, Addition, Deletion
    }

    let mut founds: Vec<Found> = Vec::new();

    let mut state = State::Other;
    let mut added = String::new();
    let mut deleted = String::new();
    let mut old_path = Path::new("");
    let mut new_path = Path::new("");

    let mut line_count: u32 = 0;
    let mut start_position: u32 = 0;

    // Read about this function in http://alexcrichton.com/git2-rs/git2/struct.Diff.html#method.print
    // It's a bit weird, but I think it will provide the necessary information.
    diff.print(DiffFormat::Patch, |delta, maybe_hunk, line| -> bool {

        // if delta.nfiles() != 2 {
        //     // This is diff with only one side, and thus can't have a move.
        //     return true;
        // }
        // assert!(delta.nfiles() == 2, "This only works on diffs between exactly 2 files. Found {} files.", delta.nfiles());

        // Thinking:
        //  * If is not a hunk, keep going.
        //  * If it's a hunk, do regex magic.
        //  * Stick regex output into a hashmap as a hash.
        //  * Later, we will iterate through and output pased on the values.
        // Filter out all the boring context lines.
        // If we're not interested in this line just return since it will iterate to the next.
        if maybe_hunk.is_none() { return true }; // Return early.

        //dump_diffline(&line);
        //dump_diffdelta(&delta);
        //dump_diffhunk(&maybe_hunk.unwrap());

        old_path = delta.old_file().path().unwrap().clone();
        new_path = delta.new_file().path().unwrap().clone();

        match line.origin() {
            // Additions
            '+' | '>' => {
                added.push_str(str::from_utf8(line.content()).unwrap());

                match state {
                    State::Deletion => {
                        founds.push(Found {
                            filename: old_path.clone(),
                            key: format_key(deleted.clone()),
                            state: FoundState::Deleted,
                            start_position: start_position,
                            line_count: line_count,
                        });
                        deleted = String::new();
                        line_count = 0;
                        start_position = line.new_lineno().unwrap();
                    },
                    State::Addition => {
                        line_count += 1;
                    },
                    State::Other => {
                        line_count = 1;
                        start_position = line.new_lineno().unwrap();
                    },
                }

                state = State::Addition;
                true
            },
            // Deletions
            '-' | '<' => {
                deleted.push_str(str::from_utf8(line.content()).unwrap());

                match state {
                    State::Addition => {
                        founds.push(Found {
                            filename: new_path.clone(),
                            key: format_key(added.clone()),
                            state: FoundState::Added,
                            start_position: start_position,
                            line_count: line_count,
                        });
                        added = String::new();
                        line_count = 0;
                        start_position = line.old_lineno().unwrap();
                    },
                    State::Deletion => {
                        line_count += 1;
                    },
                    State::Other => {
                        line_count = 1;
                        start_position = line.old_lineno().unwrap();
                    },
                }

                state = State::Deletion;
                true
            },
            // Other
            _         => {
                match state {
                    State::Addition => {
                        founds.push(Found {
                            filename: new_path.clone(),
                            key: format_key(added.clone()),
                            state: FoundState::Added,
                            start_position: start_position,
                            line_count: line_count,
                        });
                        added = String::new();
                    },
                    State::Deletion => {
                        founds.push(Found {
                            filename: old_path.clone(),
                            key: format_key(deleted.clone()),
                            state: FoundState::Deleted,
                            start_position: start_position,
                            line_count: line_count,
                        });
                        deleted = String::new();
                    },
                    _ => (),
                }

                state = State::Other;
                true
            }
        }
    });

    // Grab last one.
    match state {
        State::Addition => {
            if added.len() > 0 {
                founds.push(Found {
                    filename: new_path,
                    key: format_key(added.clone()),
                    state: FoundState::Added,
                    start_position: start_position,
                    line_count: line_count,
                });
            }
        },
        State::Deletion => {
            if deleted.len() > 0 {
                founds.push(Found {
                    filename: old_path,
                    key: format_key(deleted.clone()),
                    state: FoundState::Deleted,
                    start_position: start_position,
                    line_count: line_count,
                });
            }
        },
        _ => (),
    }

    return founds;
}

fn path_to_string(path: Path) -> String {
    String::from_utf8(path.into_vec()).unwrap()
}

fn find_moves(repo: &Repository, old: &Commit, new: &Commit) -> Result<Vec<Output>, git2::Error> {
    let old_tree = try!(old.tree());
    let new_tree = try!(new.tree());
    // Build up a diff of the two trees.
    let diff = try!(Diff::tree_to_tree(repo, Some(&old_tree), Some(&new_tree), None));

    let founds: Vec<Found> = find_additions_and_deletions(diff);

    let mut moves: Vec<Output> = Vec::new();
    let mut map: HashMap<String, Found> = HashMap::new();

    for f in founds {
        if f.key.len() == 0 { continue; }

        if map.contains_key(&f.key) {
            let q = map.get(&f.key).unwrap();

            // If both states are added, there are no moves.
            // assert!(f.state != q.state, format!("States {:?}, {:?}. Should be Addition/Deletion ", f.state, q.state));
            if (f.state == q.state) {
                return Ok(Vec::<Output>::new())
            }

            let output: Output;

            match f.state {
                FoundState::Added => {
                    output = Output {
                        old_commit: TransitOid(old.id()),
                        new_commit: TransitOid(new.id()),
                        old_filename: path_to_string(q.filename.clone()),
                        new_filename: path_to_string(f.filename.clone()),
                        origin_line: q.start_position,
                        destination_line: f.start_position,
                        num_lines: f.line_count,
                    };
                },
                FoundState::Deleted => {
                    output = Output {
                        old_commit: TransitOid(old.id()),
                        new_commit: TransitOid(new.id()),
                        old_filename: path_to_string(f.filename.clone()),
                        new_filename: path_to_string(q.filename.clone()),
                        origin_line: f.start_position,
                        destination_line: q.start_position,
                        num_lines: f.line_count,
                    };
                },
            }

            moves.push(output);
        } else {
            map.insert(f.key.clone(), f);
        }
    }

    return Ok(moves);
}

#[derive(Debug, RustcEncodable)]
struct Output {
    old_commit: TransitOid,
    new_commit: TransitOid,
    old_filename: String,
    new_filename: String,
    origin_line: u32,
    destination_line: u32,
    num_lines: u32,
}

#[derive(Debug)]
struct TransitOid(Oid);
impl Display for TransitOid {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        Display::fmt(&self.0, f)
    }
}

impl rustc_serialize::Encodable for TransitOid {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        s.emit_str(&format!("{}", self)[])
    }
}
<commit_msg>More graceful error handling.<commit_after>#![feature(core)]
#![feature(path)]

#![feature(plugin)]
#![plugin(regex_macros)]
extern crate regex;
extern crate git2;
extern crate core;
extern crate "rustc-serialize" as rustc_serialize;
extern crate docopt;

use git2::{Repository, Branch, BranchType, DiffLine,
    Commit, Diff, DiffFormat, DiffDelta, DiffHunk, Oid};
use docopt::Docopt;
use std::old_io as io;
use std::collections::HashMap;
use std::old_io::BufferedReader;
use std::old_io::File;
use std::str;
use regex::Regex;
use rustc_serialize::json::{self, ToJson, Json};
use std::fmt::Display;

// Write the Docopt usage string.
static USAGE: &'static str = "
Usage: transit <repo> [--json] [<old> <new>]

If no commits are given, transit will revwalk from latest to oldest.
";

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_repo: String,
    arg_old: Option<String>,
    arg_new: Option<String>,
    flag_json: bool,
}

fn main() {
    // Parse the args above or die.
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());
    // Pull up the revwalk.
    let path = Path::new(args.arg_repo);
    let repo = Repository::discover(&path)
        .ok().expect("Unable to find repo.");
    // If we can destructures the two optional args into real things.
    if let (Some(old_string), Some(new_string)) = (args.arg_old, args.arg_new) {
        // Compare a specific commit pair.
        // Fist, get the commits. (Error checked)
        let old = Oid::from_str(&old_string[]).and_then(|oid| repo.find_commit(oid));
        let new = Oid::from_str(&new_string[]).and_then(|oid| repo.find_commit(oid));
        if old.is_ok() && new.is_ok() {
            let output = find_moves(&repo, &old.unwrap(), &new.unwrap()).unwrap();
            make_output(output);
        } else {
            panic!("Commit ids were not valid.");
        }
    } else {
        // Revwalk.
        let mut revwalk = repo.revwalk()
            .ok().expect("Unable to get revwalk.");
        // Setup some options.
        revwalk.simplify_first_parent(); // TODO: Maybe remove?
        let mut flags = git2::Sort::empty();
        flags.insert(git2::SORT_TIME);
        flags.insert(git2::SORT_TOPOLOGICAL);
        revwalk.set_sorting(flags);
        // Push HEAD to the revwalk.
        revwalk.push_head()
            .ok().expect("Unable to push HEAD.");
        // We sadly must collect here to use `.windows()`
        let history = revwalk.filter_map(|id| repo.find_commit(id).ok())
            .collect::<Vec<Commit>>();
        // Walk through each pair of commits.
        for pair in history.windows(2) {let (old, new) = (&pair[1], &pair[0]);
            let output = find_moves(&repo, old.clone(), new.clone()).unwrap();
            if args.flag_json {
                make_json(output);
            } else {
                make_output(output);
            }
        }
    }
}

fn make_output(output: Vec<Output>) {
    println!("make_output: output.len()={}", output.len());
    for i in range(0, output.len()) {
            println!("\told_commit={}",         output[i].old_commit);
            println!("\tnew_commit={}",         output[i].new_commit);
            println!("\torigin_line={}",        output[i].origin_line);
            println!("\tdestintation_line={}",  output[i].destination_line);
            println!("\tnum_lines={}",          output[i].num_lines);
            println!("\tnew_filename={}",       output[i].new_filename);
            println!("\told_filename={}",       output[i].old_filename);
    }
}

fn make_json(output: Vec<Output>) {
    println!("{}", json::encode(&output).unwrap());
}

#[derive(Debug, PartialEq, Eq)]
enum FoundState {
    Added, Deleted
}

fn dump_diffline(line: &DiffLine) {
    // 'origin' is wrapped in pipes to ease displaying space characters.
    print!("line: old={:?} new={:?} offset={} |origin|=|{}|\n      content={}",
             line.old_lineno(), line.new_lineno(), line.content_offset(),
             line.origin(), str::from_utf8(line.content()).unwrap());
}

fn dump_diffdelta(delta: &DiffDelta) {
    println!("delta: nfiles={} status={:?} old_file=(id={} path_bytes={:?} path={:?} tsize={}) new_file=(id={} path_bytes={:?} path={:?} tsize={})",
            delta.nfiles(), delta.status(),
            delta.old_file().id(), delta.old_file().path_bytes(), delta.old_file().path(), delta.old_file().size(),
            delta.new_file().id(), delta.new_file().path_bytes(), delta.new_file().path(), delta.new_file().size());
}

fn dump_diffhunk(hunk: &DiffHunk) {
    println!("hunk: old_start={} old_lines={} new_start={} new_lines={} header={}",
            hunk.old_start(), hunk.old_lines(),
            hunk.new_start(), hunk.new_lines(),
            str::from_utf8(hunk.header()).unwrap());
}

fn format_key(key: String) -> String {
    let remove_whitespace = regex!(r"\s{2,}"); // 2 or more whitespaces    // TODO Removes whitespace from a string.
    let trim = regex!(r"^[\s]+|[\s]+$");
    let result = remove_whitespace.replace_all(key.as_slice(), "");
    trim.replace_all(result.as_slice(), "")
}

#[derive(Debug)]
struct Found {
    filename: Path,
    key: String,
    state: FoundState,
    start_position: u32,
    line_count: u32,
}

fn find_additions_and_deletions(diff: Diff) -> Vec<Found> {

    #[derive(Debug)]
    enum State {
        Other, Addition, Deletion
    }

    let mut founds: Vec<Found> = Vec::new();

    let mut state = State::Other;
    let mut added = String::new();
    let mut deleted = String::new();
    let mut old_path = Path::new("");
    let mut new_path = Path::new("");

    let mut line_count: u32 = 0;
    let mut start_position: u32 = 0;

    // Read about this function in http://alexcrichton.com/git2-rs/git2/struct.Diff.html#method.print
    // It's a bit weird, but I think it will provide the necessary information.
    diff.print(DiffFormat::Patch, |delta, maybe_hunk, line| -> bool {

        // if delta.nfiles() != 2 {
        //     // This is diff with only one side, and thus can't have a move.
        //     return true;
        // }
        // assert!(delta.nfiles() == 2, "This only works on diffs between exactly 2 files. Found {} files.", delta.nfiles());

        // Thinking:
        //  * If is not a hunk, keep going.
        //  * If it's a hunk, do regex magic.
        //  * Stick regex output into a hashmap as a hash.
        //  * Later, we will iterate through and output pased on the values.
        // Filter out all the boring context lines.
        // If we're not interested in this line just return since it will iterate to the next.
        if maybe_hunk.is_none() { return true }; // Return early.

        //dump_diffline(&line);
        //dump_diffdelta(&delta);
        //dump_diffhunk(&maybe_hunk.unwrap());

        old_path = match delta.old_file().path() {
            Some(path) => path,
            None => return false,
        };
        new_path = match delta.new_file().path() {
            Some(path) => path,
            None => return false,
        };

        match line.origin() {
            // Additions
            '+' | '>' => {
                added.push_str(str::from_utf8(line.content())
                    .ok().expect("Couldn't parse line content for added"));

                match state {
                    State::Deletion => {
                        founds.push(Found {
                            filename: old_path.clone(),
                            key: format_key(deleted.clone()),
                            state: FoundState::Deleted,
                            start_position: start_position,
                            line_count: line_count,
                        });
                        deleted = String::new();
                        line_count = 0;
                        start_position = match line.new_lineno() {
                            Some(lineno) => lineno,
                            None => return false, // Can't be a move if deleted.
                        }
                    },
                    State::Addition => {
                        line_count += 1;
                    },
                    State::Other => {
                        line_count = 1;
                        start_position = match line.new_lineno() {
                            Some(lineno) => lineno,
                            None => return false, // Can't be a move if deleted.
                        }
                    },
                }

                state = State::Addition;
                true
            },
            // Deletions
            '-' | '<' => {
                deleted.push_str(str::from_utf8(line.content())
                    .ok().expect("Couldn't parse line content for deleted"));

                match state {
                    State::Addition => {
                        founds.push(Found {
                            filename: new_path.clone(),
                            key: format_key(added.clone()),
                            state: FoundState::Added,
                            start_position: start_position,
                            line_count: line_count,
                        });
                        added = String::new();
                        line_count = 0;
                        start_position = match line.old_lineno() {
                            Some(lineno) => lineno,
                            None => return false, // Can't be a move if deleted.
                        }
                    },
                    State::Deletion => {
                        line_count += 1;
                    },
                    State::Other => {
                        line_count = 1;
                        start_position = match line.old_lineno() {
                            Some(lineno) => lineno,
                            None => return false, // Can't be a move if deleted.
                        }
                    },
                }

                state = State::Deletion;
                true
            },
            // Other
            _         => {
                match state {
                    State::Addition => {
                        founds.push(Found {
                            filename: new_path.clone(),
                            key: format_key(added.clone()),
                            state: FoundState::Added,
                            start_position: start_position,
                            line_count: line_count,
                        });
                        added = String::new();
                    },
                    State::Deletion => {
                        founds.push(Found {
                            filename: old_path.clone(),
                            key: format_key(deleted.clone()),
                            state: FoundState::Deleted,
                            start_position: start_position,
                            line_count: line_count,
                        });
                        deleted = String::new();
                    },
                    _ => (),
                }

                state = State::Other;
                true
            }
        }
    });

    // Grab last one.
    match state {
        State::Addition => {
            if added.len() > 0 {
                founds.push(Found {
                    filename: new_path,
                    key: format_key(added.clone()),
                    state: FoundState::Added,
                    start_position: start_position,
                    line_count: line_count,
                });
            }
        },
        State::Deletion => {
            if deleted.len() > 0 {
                founds.push(Found {
                    filename: old_path,
                    key: format_key(deleted.clone()),
                    state: FoundState::Deleted,
                    start_position: start_position,
                    line_count: line_count,
                });
            }
        },
        _ => (),
    }

    return founds;
}

fn path_to_string(path: Path) -> String {
    String::from_utf8(path.into_vec()).unwrap()
}

fn find_moves(repo: &Repository, old: &Commit, new: &Commit) -> Result<Vec<Output>, git2::Error> {
    let old_tree = try!(old.tree());
    let new_tree = try!(new.tree());
    // Build up a diff of the two trees.
    let diff = try!(Diff::tree_to_tree(repo, Some(&old_tree), Some(&new_tree), None));

    let founds: Vec<Found> = find_additions_and_deletions(diff);

    let mut moves: Vec<Output> = Vec::new();
    let mut map: HashMap<String, Found> = HashMap::new();

    for f in founds {
        if f.key.len() == 0 { continue; }

        if map.contains_key(&f.key) {
            let q = map.get(&f.key).unwrap();

            // If both states are added, there are no moves.
            // assert!(f.state != q.state, format!("States {:?}, {:?}. Should be Addition/Deletion ", f.state, q.state));
            if (f.state == q.state) {
                return Ok(Vec::<Output>::new())
            }

            let output: Output;

            match f.state {
                FoundState::Added => {
                    output = Output {
                        old_commit: TransitOid(old.id()),
                        new_commit: TransitOid(new.id()),
                        old_filename: path_to_string(q.filename.clone()),
                        new_filename: path_to_string(f.filename.clone()),
                        origin_line: q.start_position,
                        destination_line: f.start_position,
                        num_lines: f.line_count,
                    };
                },
                FoundState::Deleted => {
                    output = Output {
                        old_commit: TransitOid(old.id()),
                        new_commit: TransitOid(new.id()),
                        old_filename: path_to_string(f.filename.clone()),
                        new_filename: path_to_string(q.filename.clone()),
                        origin_line: f.start_position,
                        destination_line: q.start_position,
                        num_lines: f.line_count,
                    };
                },
            }

            moves.push(output);
        } else {
            map.insert(f.key.clone(), f);
        }
    }

    return Ok(moves);
}

#[derive(Debug, RustcEncodable)]
struct Output {
    old_commit: TransitOid,
    new_commit: TransitOid,
    old_filename: String,
    new_filename: String,
    origin_line: u32,
    destination_line: u32,
    num_lines: u32,
}

#[derive(Debug)]
struct TransitOid(Oid);
impl Display for TransitOid {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        Display::fmt(&self.0, f)
    }
}

impl rustc_serialize::Encodable for TransitOid {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        s.emit_str(&format!("{}", self)[])
    }
}
<|endoftext|>