<commit_before>extern crate base64;
extern crate clap;
extern crate hyper;
extern crate hyper_rustls;
extern crate pem;
#[macro_use]
extern crate prettytable;
extern crate ring;
extern crate rustls;
extern crate serde_json;
#[macro_use]
extern crate tera;

extern crate ct_tools;

use ct_tools::{crtsh, letsencrypt};
use ct_tools::common::{Log, sha256_hex};
use ct_tools::ct::submit_cert_to_logs;
use ct_tools::google::fetch_trusted_ct_logs;
use std::env;
use std::fs::File;
use std::io::{Read, Write};
use std::process::{Command, Stdio};
use std::sync::Arc;
use std::time::Duration;

fn pems_to_chain(data: &str) -> Vec<Vec<u8>> {
    pem::parse_many(data)
        .into_iter()
        .filter(|p| p.tag == "CERTIFICATE")
        .map(|p| p.contents)
        .collect()
}

fn submit(paths: clap::Values, log_urls: Option<clap::Values>) {
    let mut http_client = hyper::Client::with_connector(hyper::net::HttpsConnector::new(
        hyper_rustls::TlsClient::new(),
    ));
    http_client.set_read_timeout(Some(Duration::from_secs(15)));

    let logs = match log_urls {
        Some(urls) => {
            urls.map(|url| {
                Log {
                    url: url.to_string(),
                    description: url.to_string(),
                    is_google: false,
                }
            }).collect()
        }
        None => fetch_trusted_ct_logs(&http_client),
    };

    for path in paths {
        println!("Submitting {} ...", path);

        let mut contents = String::new();
        File::open(path)
            .unwrap()
            .read_to_string(&mut contents)
            .unwrap();

        let mut chain = pems_to_chain(&contents);
        if chain.len() == 1 {
            // TODO: There's got to be some way to do this ourselves, instead of using crt.sh as a
            // glorified AIA chaser.
            println!("Only one certificate in chain, using crt.sh to build a full chain ...");
            chain = match crtsh::build_chain_for_cert(&http_client, &chain[0]) {
                Some(c) => c,
                None => {
                    println!("Unable to build a chain");
                    continue;
                }
            }
        }
        let scts = submit_cert_to_logs(&http_client, &logs, &chain);

        if !scts.is_empty() {
            println!(
                "Find the cert on crt.sh: {}",
                crtsh::url_for_cert(&chain[0])
            );
        }
        let mut table = prettytable::Table::new();
        table.add_row(row!["Log", "SCT"]);
        for (log, sct) in scts {
            table.add_row(row![log.description, base64::encode(&sct.to_raw_bytes())]);
        }
        table.printstd();
        println!();
        println!();
    }
}

fn check(paths: clap::Values) {
    let http_client = hyper::Client::with_connector(hyper::net::HttpsConnector::new(
        hyper_rustls::TlsClient::new(),
    ));

    for path in paths {
        let mut contents = String::new();
        File::open(path)
            .unwrap()
            .read_to_string(&mut contents)
            .unwrap();

        let chain = pems_to_chain(&contents);

        let is_logged = crtsh::is_cert_logged(&http_client, &chain[0]);
        if is_logged {
            println!("{} was already logged", path);
        } else {
            println!("{} has not been logged", path);
        }
    }
}

struct HttpHandler {
    templates: tera::Tera,
    http_client: hyper::Client,
    logs: Vec<Log>,
}

impl hyper::server::Handler for HttpHandler {
    fn handle(&self, request: hyper::server::Request, response: hyper::server::Response) {
        let peer_certs = request
            .ssl::<hyper_rustls::WrappedStream>()
            .unwrap()
            .to_tls_stream()
            .get_session()
            .get_peer_certificates();

        let mut crtsh_url = None;
        if peer_certs.is_some() && request.method == hyper::method::Method::Post {
            if let Some(chain) = crtsh::build_chain_for_cert(
                &self.http_client,
                &peer_certs.as_ref().unwrap()[0].0,
            )
            {
                let scts = submit_cert_to_logs(&self.http_client, &self.logs, &chain);
                if !scts.is_empty() {
                    crtsh_url = Some(crtsh::url_for_cert(&chain[0]));
                    println!("Successfully submitted: {}", sha256_hex(&chain[0]));
                }
            }
        }

        let rendered_cert = peer_certs.map(|chain| {
            let mut process = Command::new("openssl")
                .arg("x509")
                .arg("-text")
                .arg("-noout")
                .arg("-inform")
                .arg("der")
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped())
                .spawn()
                .unwrap();
            process
                .stdin
                .as_mut()
                .unwrap()
                .write_all(&chain[0].0)
                .unwrap();
            let out = process.wait_with_output().unwrap();
            String::from_utf8_lossy(&out.stdout).into_owned()
        });

        let mut context = tera::Context::new();
        context.add("cert", &rendered_cert);
        context.add("crtsh_url", &crtsh_url);
        response
            .send(
                self.templates
                    .render("home.html", &context)
                    .unwrap()
                    .as_bytes(),
            )
            .unwrap();
    }
}

struct NoVerificationCertificateVerifier;
impl rustls::ClientCertVerifier for NoVerificationCertificateVerifier {
    fn verify_client_cert(
        &self,
        _: &rustls::RootCertStore,
        _: &[rustls::Certificate],
    ) -> Result<(), rustls::TLSError> {
        Ok(())
    }
}

fn server(local_dev: bool, domain: Option<&str>, letsencrypt_env: Option<&str>) {
    match (local_dev, domain, letsencrypt_env) {
        (true, Some(_), _) |
        (true, _, Some(_)) => {
            panic!("Can't use both --local-dev and --letsencrypt-env or --domain")
        }
        (_, Some(_), None) |
        (_, None, Some(_)) => {
            panic!("When using Let's Encrypt, must set both --letsencrypt-env and --domain")
        }
        (false, _, None) => panic!("Must set at least one of --local-dev or --letsencrypt-env"),
        _ => {}
    };

    let mut tls_config = rustls::ServerConfig::new();
    tls_config.client_auth_offer = true;
    if local_dev {
        // TODO: not all the details on the cert are perfect, but it's fine.
        let (cert, pkey) = letsencrypt::generate_temporary_cert("localhost");
        tls_config.set_single_cert(
            vec![letsencrypt::openssl_cert_to_rustls(&cert)],
            letsencrypt::openssl_pkey_to_rustls(&pkey),
        );
    } else {
        let letsencrypt_url = match letsencrypt_env.unwrap() {
            "prod" => "https://acme-v01.api.letsencrypt.org/directory",
            "dev" => "https://acme-staging.api.letsencrypt.org/directory",
            _ => unreachable!(),
        };
        let cert_cache =
            letsencrypt::DiskCache::new(env::home_dir().unwrap().join(".ct-tools").join(
                "certificates",
            ));
        tls_config.cert_resolver = Box::new(letsencrypt::AutomaticCertResolver::new(
            letsencrypt_url,
            vec![domain.unwrap().to_string()],
            cert_cache,
        ));
    }
    tls_config.set_persistence(rustls::ServerSessionMemoryCache::new(1024));
    tls_config.ticketer = rustls::Ticketer::new();
    // Disable certificate verification. In any normal context, this would be horribly insecure!
    // However, all we're doing is taking the certs and then turning around and submitting them to
    // CT logs, so it doesn't matter if they're verified.
    tls_config.dangerous().set_certificate_verifier(Box::new(
        NoVerificationCertificateVerifier,
    ));
    let tls_server = hyper_rustls::TlsServer { cfg: Arc::new(tls_config) };

    let mut http_client = hyper::Client::with_connector(hyper::net::HttpsConnector::new(
        hyper_rustls::TlsClient::new(),
    ));
    let logs = fetch_trusted_ct_logs(&http_client);
    http_client.set_write_timeout(Some(Duration::from_secs(5)));
    http_client.set_read_timeout(Some(Duration::from_secs(5)));
    let handler = HttpHandler {
        templates: compile_templates!("templates/*"),
        http_client: http_client,
        logs: logs,
    };

    let addr = if local_dev {
        "127.0.0.1:1337"
    } else {
        "0.0.0.0:443"
    };
    println!("Listening on https://{} ...", addr);
    hyper::Server::https(addr, tls_server)
        .unwrap()
        // If there aren't at least two threads, the Let's Encrypt integration will deadlock.
        .handle_threads(handler, 16)
        .unwrap();
}

fn main() {
    let matches = clap::App::new("ct-tools")
        .subcommand(
            clap::SubCommand::with_name("submit")
                .about("Directly submits certificates to CT logs")
                .arg(
                    clap::Arg::with_name("path")
                        .multiple(true)
                        .required(true)
                        .help("Path to certificate or chain"),
                )
                .arg(
                    clap::Arg::with_name("log-url")
                        .long("--log-url")
                        .takes_value(true)
                        .multiple(true)
                        .help("Log URL to submit certificate to"),
                ),
        )
        .subcommand(
            clap::SubCommand::with_name("check")
                .about("Checks whether a certificate exists in CT logs")
                .arg(
                    clap::Arg::with_name("path")
                        .multiple(true)
                        .required(true)
                        .help("Path to certificate or chain"),
                ),
        )
        .subcommand(
            clap::SubCommand::with_name("server")
                .about("Run an HTTPS server that submits client to CT logs")
                .arg(clap::Arg::with_name("local-dev").long("--local-dev").help(
                    "Local development, do not obtain a certificate",
                ))
                .arg(
                    clap::Arg::with_name("domain")
                        .takes_value(true)
                        .long("--domain")
                        .help("Domain this is running as"),
                )
                .arg(
                    clap::Arg::with_name("letsencrypt-env")
                        .takes_value(true)
                        .long("--letsencrypt-env")
                        .possible_values(&["dev", "prod"])
                        .help("Let's Encrypt environment to use"),
                ),
        )
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("submit") {
        submit(
            matches.values_of("path").unwrap(),
            matches.values_of("log-url"),
        );
    } else if let Some(matches) = matches.subcommand_matches("check") {
        check(matches.values_of("path").unwrap());
    } else if let Some(matches) = matches.subcommand_matches("server") {
        server(
            matches.is_present("local-dev"),
            matches.value_of("domain"),
            matches.value_of("letsencrypt-env"),
        );
    }
}
<commit_msg>Handle PEMs with non-utf8 data in them<commit_after>extern crate base64;
extern crate clap;
extern crate hyper;
extern crate hyper_rustls;
extern crate pem;
#[macro_use]
extern crate prettytable;
extern crate ring;
extern crate rustls;
extern crate serde_json;
#[macro_use]
extern crate tera;

extern crate ct_tools;

use ct_tools::{crtsh, letsencrypt};
use ct_tools::common::{Log, sha256_hex};
use ct_tools::ct::submit_cert_to_logs;
use ct_tools::google::fetch_trusted_ct_logs;
use std::env;
use std::fs::File;
use std::io::{Read, Write};
use std::process::{Command, Stdio};
use std::sync::Arc;
use std::time::Duration;

fn pems_to_chain(data: &[u8]) -> Vec<Vec<u8>> {
    pem::parse_many(data)
        .into_iter()
        .filter(|p| p.tag == "CERTIFICATE")
        .map(|p| p.contents)
        .collect()
}

fn submit(paths: clap::Values, log_urls: Option<clap::Values>) {
    let mut http_client = hyper::Client::with_connector(hyper::net::HttpsConnector::new(
        hyper_rustls::TlsClient::new(),
    ));
    http_client.set_read_timeout(Some(Duration::from_secs(15)));

    let logs = match log_urls {
        Some(urls) => {
            urls.map(|url| {
                Log {
                    url: url.to_string(),
                    description: url.to_string(),
                    is_google: false,
                }
            }).collect()
        }
        None => fetch_trusted_ct_logs(&http_client),
    };

    for path in paths {
        println!("Submitting {} ...", path);

        let mut contents = Vec::new();
        File::open(path)
            .unwrap()
            .read_to_end(&mut contents)
            .unwrap();

        let mut chain = pems_to_chain(&contents);
        if chain.len() == 1 {
            // TODO: There's got to be some way to do this ourselves, instead of using crt.sh as a
            // glorified AIA chaser.
            println!("Only one certificate in chain, using crt.sh to build a full chain ...");
            chain = match crtsh::build_chain_for_cert(&http_client, &chain[0]) {
                Some(c) => c,
                None => {
                    println!("Unable to build a chain");
                    continue;
                }
            }
        }
        let scts = submit_cert_to_logs(&http_client, &logs, &chain);

        if !scts.is_empty() {
            println!(
                "Find the cert on crt.sh: {}",
                crtsh::url_for_cert(&chain[0])
            );
        }
        let mut table = prettytable::Table::new();
        table.add_row(row!["Log", "SCT"]);
        for (log, sct) in scts {
            table.add_row(row![log.description, base64::encode(&sct.to_raw_bytes())]);
        }
        table.printstd();
        println!();
        println!();
    }
}

fn check(paths: clap::Values) {
    let http_client = hyper::Client::with_connector(hyper::net::HttpsConnector::new(
        hyper_rustls::TlsClient::new(),
    ));

    for path in paths {
        let mut contents = Vec::new();
        File::open(path)
            .unwrap()
            .read_to_end(&mut contents)
            .unwrap();

        let chain = pems_to_chain(&contents);

        let is_logged = crtsh::is_cert_logged(&http_client, &chain[0]);
        if is_logged {
            println!("{} was already logged", path);
        } else {
            println!("{} has not been logged", path);
        }
    }
}

struct HttpHandler {
    templates: tera::Tera,
    http_client: hyper::Client,
    logs: Vec<Log>,
}

impl hyper::server::Handler for HttpHandler {
    fn handle(&self, request: hyper::server::Request, response: hyper::server::Response) {
        let peer_certs = request
            .ssl::<hyper_rustls::WrappedStream>()
            .unwrap()
            .to_tls_stream()
            .get_session()
            .get_peer_certificates();

        let mut crtsh_url = None;
        if peer_certs.is_some() && request.method == hyper::method::Method::Post {
            if let Some(chain) = crtsh::build_chain_for_cert(
                &self.http_client,
                &peer_certs.as_ref().unwrap()[0].0,
            )
            {
                let scts = submit_cert_to_logs(&self.http_client, &self.logs, &chain);
                if !scts.is_empty() {
                    crtsh_url = Some(crtsh::url_for_cert(&chain[0]));
                    println!("Successfully submitted: {}", sha256_hex(&chain[0]));
                }
            }
        }

        let rendered_cert = peer_certs.map(|chain| {
            let mut process = Command::new("openssl")
                .arg("x509")
                .arg("-text")
                .arg("-noout")
                .arg("-inform")
                .arg("der")
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped())
                .spawn()
                .unwrap();
            process
                .stdin
                .as_mut()
                .unwrap()
                .write_all(&chain[0].0)
                .unwrap();
            let out = process.wait_with_output().unwrap();
            String::from_utf8_lossy(&out.stdout).into_owned()
        });

        let mut context = tera::Context::new();
        context.add("cert", &rendered_cert);
        context.add("crtsh_url", &crtsh_url);
        response
            .send(
                self.templates
                    .render("home.html", &context)
                    .unwrap()
                    .as_bytes(),
            )
            .unwrap();
    }
}

struct NoVerificationCertificateVerifier;
impl rustls::ClientCertVerifier for NoVerificationCertificateVerifier {
    fn verify_client_cert(
        &self,
        _: &rustls::RootCertStore,
        _: &[rustls::Certificate],
    ) -> Result<(), rustls::TLSError> {
        Ok(())
    }
}

fn server(local_dev: bool, domain: Option<&str>, letsencrypt_env: Option<&str>) {
    match (local_dev, domain, letsencrypt_env) {
        (true, Some(_), _) |
        (true, _, Some(_)) => {
            panic!("Can't use both --local-dev and --letsencrypt-env or --domain")
        }
        (_, Some(_), None) |
        (_, None, Some(_)) => {
            panic!("When using Let's Encrypt, must set both --letsencrypt-env and --domain")
        }
        (false, _, None) => panic!("Must set at least one of --local-dev or --letsencrypt-env"),
        _ => {}
    };

    let mut tls_config = rustls::ServerConfig::new();
    tls_config.client_auth_offer = true;
    if local_dev {
        // TODO: not all the details on the cert are perfect, but it's fine.
        let (cert, pkey) = letsencrypt::generate_temporary_cert("localhost");
        tls_config.set_single_cert(
            vec![letsencrypt::openssl_cert_to_rustls(&cert)],
            letsencrypt::openssl_pkey_to_rustls(&pkey),
        );
    } else {
        let letsencrypt_url = match letsencrypt_env.unwrap() {
            "prod" => "https://acme-v01.api.letsencrypt.org/directory",
            "dev" => "https://acme-staging.api.letsencrypt.org/directory",
            _ => unreachable!(),
        };
        let cert_cache =
            letsencrypt::DiskCache::new(env::home_dir().unwrap().join(".ct-tools").join(
                "certificates",
            ));
        tls_config.cert_resolver = Box::new(letsencrypt::AutomaticCertResolver::new(
            letsencrypt_url,
            vec![domain.unwrap().to_string()],
            cert_cache,
        ));
    }
    tls_config.set_persistence(rustls::ServerSessionMemoryCache::new(1024));
    tls_config.ticketer = rustls::Ticketer::new();
    // Disable certificate verification. In any normal context, this would be horribly insecure!
    // However, all we're doing is taking the certs and then turning around and submitting them to
    // CT logs, so it doesn't matter if they're verified.
    tls_config.dangerous().set_certificate_verifier(Box::new(
        NoVerificationCertificateVerifier,
    ));
    let tls_server = hyper_rustls::TlsServer { cfg: Arc::new(tls_config) };

    let mut http_client = hyper::Client::with_connector(hyper::net::HttpsConnector::new(
        hyper_rustls::TlsClient::new(),
    ));
    let logs = fetch_trusted_ct_logs(&http_client);
    http_client.set_write_timeout(Some(Duration::from_secs(5)));
    http_client.set_read_timeout(Some(Duration::from_secs(5)));
    let handler = HttpHandler {
        templates: compile_templates!("templates/*"),
        http_client: http_client,
        logs: logs,
    };

    let addr = if local_dev {
        "127.0.0.1:1337"
    } else {
        "0.0.0.0:443"
    };
    println!("Listening on https://{} ...", addr);
    hyper::Server::https(addr, tls_server)
        .unwrap()
        // If there aren't at least two threads, the Let's Encrypt integration will deadlock.
        .handle_threads(handler, 16)
        .unwrap();
}

fn main() {
    let matches = clap::App::new("ct-tools")
        .subcommand(
            clap::SubCommand::with_name("submit")
                .about("Directly submits certificates to CT logs")
                .arg(
                    clap::Arg::with_name("path")
                        .multiple(true)
                        .required(true)
                        .help("Path to certificate or chain"),
                )
                .arg(
                    clap::Arg::with_name("log-url")
                        .long("--log-url")
                        .takes_value(true)
                        .multiple(true)
                        .help("Log URL to submit certificate to"),
                ),
        )
        .subcommand(
            clap::SubCommand::with_name("check")
                .about("Checks whether a certificate exists in CT logs")
                .arg(
                    clap::Arg::with_name("path")
                        .multiple(true)
                        .required(true)
                        .help("Path to certificate or chain"),
                ),
        )
        .subcommand(
            clap::SubCommand::with_name("server")
                .about("Run an HTTPS server that submits client to CT logs")
                .arg(clap::Arg::with_name("local-dev").long("--local-dev").help(
                    "Local development, do not obtain a certificate",
                ))
                .arg(
                    clap::Arg::with_name("domain")
                        .takes_value(true)
                        .long("--domain")
                        .help("Domain this is running as"),
                )
                .arg(
                    clap::Arg::with_name("letsencrypt-env")
                        .takes_value(true)
                        .long("--letsencrypt-env")
                        .possible_values(&["dev", "prod"])
                        .help("Let's Encrypt environment to use"),
                ),
        )
        .get_matches();

    if let Some(matches) = matches.subcommand_matches("submit") {
        submit(
            matches.values_of("path").unwrap(),
            matches.values_of("log-url"),
        );
    } else if let Some(matches) = matches.subcommand_matches("check") {
        check(matches.values_of("path").unwrap());
    } else if let Some(matches) = matches.subcommand_matches("server") {
        server(
            matches.is_present("local-dev"),
            matches.value_of("domain"),
            matches.value_of("letsencrypt-env"),
        );
    }
}
<|endoftext|>


<commit_before>extern crate memmap;
extern crate thread_scoped;
extern crate libc;
extern crate memchr;

use std::env;
use std::thread;
use std::io::{BufReader, BufRead, Read};
use std::fs::File;
use std::process::exit;
use memmap::{Mmap, Protection};
use std::str;
use libc::{madvise, POSIX_MADV_SEQUENTIAL, MADV_SEQUENTIAL, MADV_WILLNEED, posix_fadvise,
           POSIX_FADV_SEQUENTIAL};
use std::os::unix::io::AsRawFd;
use memchr::memchr;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("Not enough args!");
        exit(1);
    }
    let filepath = &args[1];
    // let total_lines = count_bufread_serial(filepath);
    // let total_lines = count_bufread_serial_fadvise(filepath);
    // let total_lines = count_manual_read_fadvise(filepath);
    // let total_lines = count_manual_read_parallel_fadvise_memchr(filepath);
    // let total_lines = count_manual_read_memchr_fadvise(filepath);
    // let total_lines = count_mmap_serial(filepath);
    // let total_lines = count_mmap_serial_memchr(filepath);
    // let total_lines = count_mmap_serial_madvise_memchr(filepath);
    // let total_lines = count_mmap_parallel(filepath);
    // let total_lines = count_mmap_parallel_memchr(filepath);
    let total_lines = count_mmap_parallel_memchr_populate(filepath);
    println!("lines: {}", total_lines);
}

// NOTE(cgag): Super slow!  Why?  Takes 6 seconds.  More than twice as slow as
// serial mmap.
fn count_bufread_serial(filepath: &str) -> u64 {
    let file = File::open(filepath).unwrap();
    let reader = BufReader::new(file);

    reader.lines().fold(0, |acc, _| acc + 1)
}

fn count_bufread_serial_fadvise(filepath: &str) -> u64 {
    let file = File::open(filepath).unwrap();
    let file_size: i64 = file.metadata().unwrap().len() as i64;
    let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, file_size, POSIX_FADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in fadvise: {}", ret);
        exit(ret);
    }
    let reader = BufReader::new(file);

    reader.lines().fold(0, |acc, _| acc + 1)
}

fn count_manual_read_fadvise(filepath: &str) -> u64 {
    let mut file = File::open(filepath).unwrap();
    let file_size: i64 = file.metadata().unwrap().len() as i64;
    // let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, file_size, POSIX_FADV_SEQUENTIAL) };
    let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, 0, POSIX_FADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in fadvise: {}", ret);
        exit(ret);
    }

    // TODO(cgag): what's a page size? how many pages should this be?  Should
    // it be the size of L3 cache, L2, etc?
    // Using 4096 as oit's listed as the cache size in /proc/cpuinfo
    let mut buf: [u8; 1024 * 16] = [0; 1024 * 16];
    let mut lines = 0;
    loop {
        match file.read(&mut buf) {
            Ok(0) => break,
            Ok(n) => {
                for i in 0..n {
                    if buf[i] == b'\n' {
                        lines += 1;
                    }
                }
            }
            _ => panic!("shit"),
        }
    }
    lines
}

// TODO(cgag): read more about how memchr is implemented
fn count_manual_read_memchr_fadvise(filepath: &str) -> u64 {
    let mut file = File::open(filepath).unwrap();
    let file_size: i64 = file.metadata().unwrap().len() as i64;
    // let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, file_size, POSIX_FADV_SEQUENTIAL) };
    let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, 0, POSIX_FADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in fadvise: {}", ret);
        exit(ret);
    }

    // TODO(cgag): what's a page size? how many pages should this be?  Should
    // it be the size of L3 cache, L2, etc?
    // Using 4096 as oit's listed as the cache size in /proc/cpuinfo
    let mut buf: [u8; 1024 * 16] = [0; 1024 * 16];
    let mut lines = 0;
    loop {
        match file.read(&mut buf) {
            Ok(0) => break,
            Ok(n) => lines += count_buf_lines(&buf[0..n]),
            _ => panic!("shit"),
        }
    }
    lines
}

fn count_mmap_serial(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };

    let mut lines = 0;
    for byte in bytes {
        if *byte == b'\n' {
            lines += 1;
        }
    }
    lines
}

fn count_mmap_serial_memchr(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };
    count_buf_lines(bytes)
}

fn count_mmap_serial_madvise(filepath: &str) -> u64 {
    let mut fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let mut bytes: &mut [u8] = unsafe { fmmap.as_mut_slice() };

    // TODO(cgag): try MAP_POPULATE?
    let mut bytes_ptr = &mut *bytes as *mut _ as *mut libc::c_void;
    let ret = unsafe { madvise(bytes_ptr, bytes.len(), POSIX_MADV_SEQUENTIAL) };
    // let ret = unsafe { madvise(bytes_ptr, bytes.len(), MADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in madvise: {}", ret);
        exit(ret);
    }

    let mut lines = 0;
    for byte in bytes {
        if *byte == b'\n' {
            lines += 1;
        }
    }
    lines
}

fn count_mmap_serial_madvise_memchr(filepath: &str) -> u64 {
    let mut fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let mut bytes: &mut [u8] = unsafe { fmmap.as_mut_slice() };

    // TODO(cgag): try MAP_POPULATE?
    let mut bytes_ptr = &mut *bytes as *mut _ as *mut libc::c_void;
    let ret = unsafe { madvise(bytes_ptr, bytes.len(), POSIX_MADV_SEQUENTIAL) };
    // let ret = unsafe { madvise(bytes_ptr, bytes.len(), MADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in madvise: {}", ret);
        exit(ret);
    }

    count_buf_lines(bytes)
}

fn count_mmap_parallel(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };

    let mut handles: Vec<thread_scoped::JoinGuard<u64>> = Vec::new();

    for chunk in bytes.chunks(bytes.len() / 4) {
        unsafe {
            let t = thread_scoped::scoped(move || {
                let mut lines = 0;
                for byte in chunk {
                    if *byte == b'\n' {
                        lines += 1;
                    }
                }
                lines
            });
            handles.push(t);
        };
    }

    let mut total_lines = 0;
    for h in handles {
        total_lines += h.join()
    }
    total_lines
}

fn count_mmap_parallel_memchr(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };

    let mut handles: Vec<thread_scoped::JoinGuard<u64>> = Vec::new();

    for chunk in bytes.chunks(bytes.len() / 4) {
        unsafe {
            let t = thread_scoped::scoped(move || count_buf_lines(chunk));
            handles.push(t);
        };
    }

    let mut total_lines = 0;
    for h in handles {
        total_lines += h.join()
    }
    total_lines
}

fn count_mmap_parallel_memchr_populate(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };

    let mut handles: Vec<thread_scoped::JoinGuard<u64>> = Vec::new();

    for chunk in bytes.chunks(bytes.len() / 4) {
        unsafe {
            let t = thread_scoped::scoped(move || count_buf_lines(chunk));
            handles.push(t);
        };
    }

    let mut total_lines = 0;
    for h in handles {
        total_lines += h.join()
    }
    total_lines
}

fn count_buf_lines(buf: &[u8]) -> u64 {
    let mut lines = 0;
    let mut start = 0;
    loop {
        match memchr(b'\n', &buf[start..buf.len()]) {
            Some(nn) => {
                start = start + nn + 1;
                lines += 1;
            }
            None => {
                break;
            }
        }
    }
    lines
}
<commit_msg>diff main<commit_after>extern crate memmap;
extern crate thread_scoped;
extern crate libc;
extern crate memchr;

use std::env;
use std::thread;
use std::io::{BufReader, BufRead, Read};
use std::fs::File;
use std::process::exit;
use memmap::{Mmap, Protection};
use std::str;
use libc::{madvise, POSIX_MADV_SEQUENTIAL, MADV_SEQUENTIAL, MADV_WILLNEED, posix_fadvise,
           POSIX_FADV_SEQUENTIAL};
use std::os::unix::io::AsRawFd;
use memchr::memchr;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("Not enough args!");
        exit(1);
    }
    let filepath = &args[1];
    // let total_lines = count_bufread_serial(filepath);
    // let total_lines = count_bufread_serial_fadvise(filepath);
    // let total_lines = count_manual_read_fadvise(filepath);
    // let total_lines = count_manual_read_parallel_fadvise_memchr(filepath);
    // let total_lines = count_manual_read_memchr_fadvise(filepath);
    // let total_lines = count_mmap_serial(filepath);
    // let total_lines = count_mmap_serial_memchr(filepath);
    // let total_lines = count_mmap_serial_madvise_memchr(filepath);
    // let total_lines = count_mmap_parallel(filepath);
    let total_lines = count_mmap_parallel_memchr(filepath);
    println!("lines: {}", total_lines);
}

// NOTE(cgag): Super slow!  Why?  Takes 6 seconds.  More than twice as slow as
// serial mmap.
fn count_bufread_serial(filepath: &str) -> u64 {
    let file = File::open(filepath).unwrap();
    let reader = BufReader::new(file);

    reader.lines().fold(0, |acc, _| acc + 1)
}

fn count_bufread_serial_fadvise(filepath: &str) -> u64 {
    let file = File::open(filepath).unwrap();
    let file_size: i64 = file.metadata().unwrap().len() as i64;
    let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, file_size, POSIX_FADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in fadvise: {}", ret);
        exit(ret);
    }
    let reader = BufReader::new(file);

    reader.lines().fold(0, |acc, _| acc + 1)
}

fn count_manual_read_fadvise(filepath: &str) -> u64 {
    let mut file = File::open(filepath).unwrap();
    let file_size: i64 = file.metadata().unwrap().len() as i64;
    // let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, file_size, POSIX_FADV_SEQUENTIAL) };
    let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, 0, POSIX_FADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in fadvise: {}", ret);
        exit(ret);
    }

    // TODO(cgag): what's a page size? how many pages should this be?  Should
    // it be the size of L3 cache, L2, etc?
    // Using 4096 as oit's listed as the cache size in /proc/cpuinfo
    let mut buf: [u8; 1024 * 16] = [0; 1024 * 16];
    let mut lines = 0;
    loop {
        match file.read(&mut buf) {
            Ok(0) => break,
            Ok(n) => {
                for i in 0..n {
                    if buf[i] == b'\n' {
                        lines += 1;
                    }
                }
            }
            _ => panic!("shit"),
        }
    }
    lines
}

// TODO(cgag): read more about how memchr is implemented
fn count_manual_read_memchr_fadvise(filepath: &str) -> u64 {
    let mut file = File::open(filepath).unwrap();
    let file_size: i64 = file.metadata().unwrap().len() as i64;
    // let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, file_size, POSIX_FADV_SEQUENTIAL) };
    let ret = unsafe { posix_fadvise(file.as_raw_fd(), 0, 0, POSIX_FADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in fadvise: {}", ret);
        exit(ret);
    }

    // TODO(cgag): what's a page size? how many pages should this be?  Should
    // it be the size of L3 cache, L2, etc?
    // Using 4096 as oit's listed as the cache size in /proc/cpuinfo
    let mut buf: [u8; 1024 * 16] = [0; 1024 * 16];
    let mut lines = 0;
    loop {
        match file.read(&mut buf) {
            Ok(0) => break,
            Ok(n) => lines += count_buf_lines(&buf[0..n]),
            _ => panic!("shit"),
        }
    }
    lines
}

fn count_mmap_serial(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };

    let mut lines = 0;
    for byte in bytes {
        if *byte == b'\n' {
            lines += 1;
        }
    }
    lines
}

fn count_mmap_serial_memchr(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };
    count_buf_lines(bytes)
}

fn count_mmap_serial_madvise(filepath: &str) -> u64 {
    let mut fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let mut bytes: &mut [u8] = unsafe { fmmap.as_mut_slice() };

    // TODO(cgag): try MAP_POPULATE?
    let mut bytes_ptr = &mut *bytes as *mut _ as *mut libc::c_void;
    let ret = unsafe { madvise(bytes_ptr, bytes.len(), POSIX_MADV_SEQUENTIAL) };
    // let ret = unsafe { madvise(bytes_ptr, bytes.len(), MADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in madvise: {}", ret);
        exit(ret);
    }

    let mut lines = 0;
    for byte in bytes {
        if *byte == b'\n' {
            lines += 1;
        }
    }
    lines
}

fn count_mmap_serial_madvise_memchr(filepath: &str) -> u64 {
    let mut fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let mut bytes: &mut [u8] = unsafe { fmmap.as_mut_slice() };

    // TODO(cgag): try MAP_POPULATE?
    let mut bytes_ptr = &mut *bytes as *mut _ as *mut libc::c_void;
    let ret = unsafe { madvise(bytes_ptr, bytes.len(), POSIX_MADV_SEQUENTIAL) };
    // let ret = unsafe { madvise(bytes_ptr, bytes.len(), MADV_SEQUENTIAL) };
    if ret != 0 {
        println!("error in madvise: {}", ret);
        exit(ret);
    }

    count_buf_lines(bytes)
}

fn count_mmap_parallel(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };

    let mut handles: Vec<thread_scoped::JoinGuard<u64>> = Vec::new();

    for chunk in bytes.chunks(bytes.len() / 4) {
        unsafe {
            let t = thread_scoped::scoped(move || {
                let mut lines = 0;
                for byte in chunk {
                    if *byte == b'\n' {
                        lines += 1;
                    }
                }
                lines
            });
            handles.push(t);
        };
    }

    let mut total_lines = 0;
    for h in handles {
        total_lines += h.join()
    }
    total_lines
}

fn count_mmap_parallel_memchr(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };

    let mut handles: Vec<thread_scoped::JoinGuard<u64>> = Vec::new();

    for chunk in bytes.chunks(bytes.len() / 4) {
        unsafe {
            let t = thread_scoped::scoped(move || count_buf_lines(chunk));
            handles.push(t);
        };
    }

    let mut total_lines = 0;
    for h in handles {
        total_lines += h.join()
    }
    total_lines
}

fn count_mmap_parallel_memchr_populate(filepath: &str) -> u64 {
    let fmmap = Mmap::open_path(filepath, Protection::Read).expect("mmap err");
    let bytes: &[u8] = unsafe { fmmap.as_slice() };

    let mut handles: Vec<thread_scoped::JoinGuard<u64>> = Vec::new();

    for chunk in bytes.chunks(bytes.len() / 4) {
        unsafe {
            let t = thread_scoped::scoped(move || count_buf_lines(chunk));
            handles.push(t);
        };
    }

    let mut total_lines = 0;
    for h in handles {
        total_lines += h.join()
    }
    total_lines
}

fn count_buf_lines(buf: &[u8]) -> u64 {
    let mut lines = 0;
    let mut start = 0;
    loop {
        match memchr(b'\n', &buf[start..buf.len()]) {
            Some(nn) => {
                start = start + nn + 1;
                lines += 1;
            }
            None => {
                break;
            }
        }
    }
    lines
}
<|endoftext|>


<commit_before>//#![feature(convert)]

#[macro_use]
extern crate nickel;
extern crate rustc_serialize;
extern crate hyper;
extern crate nickel_postgres;
extern crate postgres;
extern crate r2d2_postgres;
extern crate r2d2;
extern crate time;

use nickel::{Nickel, HttpRouter, QueryString};
use nickel::extensions::Redirect;
use nickel::status::StatusCode;
use rustc_serialize::json;
use hyper::Client;
use hyper::header::Connection;
use hyper::status::StatusClass;
use hyper::client::response::Response;
use std::io::Read;
use std::env;
use r2d2::{NopErrorHandler, PooledConnection};
use r2d2_postgres::{TlsMode, PostgresConnectionManager};
use nickel_postgres::{PostgresMiddleware, PostgresRequestExtensions};
use time::now;

#[derive(RustcDecodable, Debug)]
struct Data {
    data: User,
}

#[derive(RustcDecodable, Debug)]
struct User {
    id: i32,
    email: String,
    created_at: String,
    updated_at: String,
    first_name: String,
    last_name: String,
    major: String,
    shirt_size: String,
    dietary_restrictions: String,
    special_needs: Option<String>,
    date_of_birth: String,
    gender: String,
    phone_number: String,
    school: School,
}

#[derive(RustcDecodable, Debug)]
struct School {
    id: i32,
    name: String,
}

#[derive(RustcDecodable, Debug)]
struct TokenResp {
    access_token: String,
    token_type: String,
    scope: String,
    created_at: u64,
}

#[derive(RustcEncodable, Debug)]
struct Payload {
    channel: String,
    username: String,
    text: String,
    icon_emoji: String,
}

#[derive(Debug)]
enum ApiError {
    ClientError,
    ServerError,
}

#[derive(Debug)]
enum RequestError {
    Hyper(hyper::error::Error),
    Io(std::io::Error),
    JsonEnc(rustc_serialize::json::EncoderError),
    JsonDec(rustc_serialize::json::DecoderError),
    Api(ApiError),
}

// `env_err!(var, message)` gets `var` from the environment, and panics with
// `message` if that fails
// `env_err!(var)` gets `var` from the environment and panics with `"Failed to
// get env var $var"` if that fails
macro_rules! env_err {
    ( $var:expr, $error:expr ) => {
        env::var($var).expect($error)
    };
    ( $var:expr ) => {
        env_err!($var, &format!("Failed to get env var {}", $var))
    }
}

// Error handling code.  Expands to impl-ing From for enum_t from from_t to to_t
// See https://doc.rust-lang.org/book/error-handling.html for a full description
macro_rules! impl_from {
    ( $from_t:path, $to_t:path, $enum_t:path ) => {
        impl From<$from_t> for $enum_t {
            fn from(err: $from_t) -> $enum_t {
                $to_t(err)
            }
        }
    }
}

// All the impls for RequestError
impl_from!(hyper::error::Error, RequestError::Hyper, RequestError);
impl_from!(std::io::Error, RequestError::Io, RequestError);
impl_from!(rustc_serialize::json::DecoderError,
           RequestError::JsonDec,
           RequestError);
impl_from!(rustc_serialize::json::EncoderError,
           RequestError::JsonEnc,
           RequestError);
impl_from!(ApiError, RequestError::Api, RequestError);

// Checks for 4xx or 5xx errors and returns the appropriate ApiError
fn check_http_error(res: &Response) -> Result<(), ApiError> {
    match res.status.class() {
        StatusClass::ClientError => Err(ApiError::ClientError),
        StatusClass::ServerError => Err(ApiError::ServerError),
        _ => Ok(()),
    }
}

// Do the OAUTH stuff for my.mlh.io
// TODO: Something better about the long api urls
fn do_request(code: &str) -> Result<Data, RequestError> {
    let id = env_err!("ID");
    let secret = env_err!("SECRET");
    let redirect = env_err!("REDIRECT");

    let url = format!("https://my.mlh.io/oauth/token?client_id={}&\
                      client_secret={}&code={}&redirect_uri={}&\
                      grant_type=authorization_code",
                      id,
                      secret,
                      code,
                      redirect);

    let client = Client::new();
    let mut res = try!(client.post(&url)
        .header(Connection::close())
        .send());

    try!(check_http_error(&res));

    let mut body = String::new();
    try!(res.read_to_string(&mut body));

    let token: TokenResp = try!(json::decode(&body));
    let url2 = format!("https://my.mlh.io/api/v2/user?access_token={}",
                       token.access_token);

    res = try!(client.get(&url2)
        .header(Connection::close())
        .send());

    try!(check_http_error(&res));

    body = String::new();
    try!(res.read_to_string(&mut body));

    let person_data: Data = try!(json::decode(&body));

    Ok(person_data)
}

// Send a message to slack when a new user signs up
fn slack_send(user: User) -> Result<(), RequestError> {
    let url = env_err!("SLACKURL");
    let client = Client::new();
    let payload = Payload {
        channel: "#signups".to_string(),
        username: "Signup bot".to_string(),
        icon_emoji: ":hackcwru:".to_string(),
        text: format!("{}, a {} major from {}, has signed up!",
                      user.first_name,
                      user.major,
                      user.school.name),
    };
    let payload_str = try!(json::encode(&payload));
    let res = try!(client.post(&url)
        .body(&payload_str)
        .send());

    try!(check_http_error(&res));

    Ok(())
}

// Create the table if it doesn't exist.  Runs on each startup
fn create_table(conn: PooledConnection<PostgresConnectionManager>) {
    let _r =
        conn.execute("CREATE TABLE IF NOT EXISTS person (
                id SERIAL PRIMARY KEY,
                email VARCHAR NOT NULL,
                created_at VARCHAR NOT NULL,
                updated_at VARCHAR NOT NULL,
                first_name VARCHAR NOT NULL,
                last_name VARCHAR NOT NULL,
                graduation VARCHAR,
                major VARCHAR NOT NULL,
                shirt_size VARCHAR NOT NULL,
                dietary_restrictions VARCHAR NOT NULL,
                special_needs VARCHAR,
                date_of_birth VARCHAR NOT NULL,
                gender VARCHAR NOT NULL,
                phone_number VARCHAR NOT NULL,
                school_id integer,
                school_name VARCHAR
                )",
            &[]
        );
    let _r = conn.execute("ALTER TABLE person ADD COLUMN year integer \
                          NOT NULL DEFAULT 2016",
                          &[]);
    let _r = conn.execute("ALTER TABLE person ALTER COLUMN graduation \
                          DROP NOT NULL",
                          &[]);
    let _r = conn.execute(
            "ALTER TABLE person ADD COLUMN signup_time VARCHAR", &[]
        );
    let _r = conn.execute(
            "ALTER TABLE person ADD COLUMN signed_in
            BOOLEAN NOT NULL DEFAULT FALSE",
        &[]
        );
    let _r = conn.execute(
            "ALTER TABLE person ADD COLUMN signed_in_time VARCHAR", &[]
        );

    let _r = conn.execute(
            "ALTER TABLE person DROP CONSTRAINT person_pkey", &[]
        );

    let _r = conn.execute(
            "ALTER TABLE person ADD PRIMARY KEY (id, year)", &[]
        );
}

fn main() {
    let mut app = Nickel::new();

    let postgres_url = env_err!("DATABASE");
    let dbpool =
        PostgresMiddleware::new(&*postgres_url, TlsMode::None, 5, Box::new(NopErrorHandler))
            .expect("Failed to start PostgresMiddleware");

    create_table(dbpool.pool.clone().get().unwrap());
    app.utilize(dbpool);

    app.get("/callback", middleware! { |request, mut response|
        let conn = request.db_conn();
        let redirect_success = env_err!("FINALREDIRECT");
        let redirect_failure = env_err!("REDIRECTERR");
        let year = env_err!("YEAR").parse::<i32>().unwrap();
        let code = match request.query().get("code") {
            Some(s) => s,
            None => {
                println!("No code, quitting");
                return response.error(StatusCode::BadRequest, "Failed");
            }
        };
        let user_data = match do_request(code) {
            Ok(s) => s,
            Err(e) => {
                println!("Get user data failed with error: {:?}", e);
                return response
                    .error(StatusCode::InternalServerError, "Failed");
            }
        }.data;

        let now = now();
        let now_s = format!("{}", now.rfc3339());
        // TODO: make this a function
        let r = conn.execute(
                "INSERT INTO person (id, email, created_at, updated_at,
                first_name, last_name, major, shirt_size,
                dietary_restrictions, special_needs, date_of_birth, gender,
                phone_number, school_id, school_name, year, signup_time) VALUES
                ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14,
                $15, $16, $17)",
                &[&user_data.id, &user_data.email, &user_data.created_at,
                  &user_data.updated_at, &user_data.first_name,
                  &user_data.last_name, &user_data.major,
                  &user_data.shirt_size, &user_data.dietary_restrictions,
                  &user_data.special_needs, &user_data.date_of_birth,
                  &user_data.gender, &user_data.phone_number,
                  &user_data.school.id, &user_data.school.name, &year, &now_s]
            );
        let redirect = match r {
            Ok(v) => {
                println!("Add to database succeeded with status {:?}", v);
                match slack_send(user_data) {
                    Ok(_) => println!("Slack send worked"),
                    Err(e) => println!("Slack send failed with error: {:?}", e)
                }
                redirect_success
            },
            Err(e) => {
                println!("Add to database failed: {:?}", e);
                redirect_failure
            }
        };

        return response.redirect(&redirect[..]);
    });

    app.get("/start", middleware! { |_req, response|
        let id = env_err!("ID");
        let redirect = env_err!("REDIRECT");
        return response.redirect(
            format!(
                    "http://my.mlh.io/oauth/authorize?client_id={}&\
                    redirect_uri={}&response_type=code",
                    id, redirect
                )
            );
    });

    let bind = match env::var("BIND") {
        Ok(v) => v,
        Err(_) => "127.0.0.1:8080".to_string(),
    };

    let _ = app.listen(&bind[..]);
}
<commit_msg>Email page<commit_after>//#![feature(convert)]

#[macro_use]
extern crate nickel;
extern crate rustc_serialize;
extern crate hyper;
extern crate nickel_postgres;
extern crate postgres;
extern crate r2d2_postgres;
extern crate r2d2;
extern crate time;

use nickel::{Nickel, HttpRouter, QueryString};
use nickel::extensions::Redirect;
use nickel::status::StatusCode;
use rustc_serialize::json;
use hyper::Client;
use hyper::header::Connection;
use hyper::status::StatusClass;
use hyper::client::response::Response;
use std::io::Read;
use std::env;
use r2d2::{NopErrorHandler, PooledConnection};
use r2d2_postgres::{TlsMode, PostgresConnectionManager};
use nickel_postgres::{PostgresMiddleware, PostgresRequestExtensions};
use time::now;

#[derive(RustcDecodable, Debug)]
struct Data {
    data: User,
}

#[derive(RustcDecodable, Debug)]
struct User {
    id: i32,
    email: String,
    created_at: String,
    updated_at: String,
    first_name: String,
    last_name: String,
    major: String,
    shirt_size: String,
    dietary_restrictions: String,
    special_needs: Option<String>,
    date_of_birth: String,
    gender: String,
    phone_number: String,
    school: School,
}

#[derive(RustcDecodable, Debug)]
struct School {
    id: i32,
    name: String,
}

#[derive(RustcDecodable, Debug)]
struct TokenResp {
    access_token: String,
    token_type: String,
    scope: String,
    created_at: u64,
}

#[derive(RustcEncodable, Debug)]
struct Payload {
    channel: String,
    username: String,
    text: String,
    icon_emoji: String,
}

#[derive(Debug)]
enum ApiError {
    ClientError,
    ServerError,
}

#[derive(Debug)]
enum RequestError {
    Hyper(hyper::error::Error),
    Io(std::io::Error),
    JsonEnc(rustc_serialize::json::EncoderError),
    JsonDec(rustc_serialize::json::DecoderError),
    Api(ApiError),
}

// `env_err!(var, message)` gets `var` from the environment, and panics with
// `message` if that fails
// `env_err!(var)` gets `var` from the environment and panics with `"Failed to
// get env var $var"` if that fails
macro_rules! env_err {
    ( $var:expr, $error:expr ) => {
        env::var($var).expect($error)
    };
    ( $var:expr ) => {
        env_err!($var, &format!("Failed to get env var {}", $var))
    }
}

// Error handling code.  Expands to impl-ing From for enum_t from from_t to to_t
// See https://doc.rust-lang.org/book/error-handling.html for a full description
macro_rules! impl_from {
    ( $from_t:path, $to_t:path, $enum_t:path ) => {
        impl From<$from_t> for $enum_t {
            fn from(err: $from_t) -> $enum_t {
                $to_t(err)
            }
        }
    }
}

// All the impls for RequestError
impl_from!(hyper::error::Error, RequestError::Hyper, RequestError);
impl_from!(std::io::Error, RequestError::Io, RequestError);
impl_from!(rustc_serialize::json::DecoderError,
           RequestError::JsonDec,
           RequestError);
impl_from!(rustc_serialize::json::EncoderError,
           RequestError::JsonEnc,
           RequestError);
impl_from!(ApiError, RequestError::Api, RequestError);

// Checks for 4xx or 5xx errors and returns the appropriate ApiError
fn check_http_error(res: &Response) -> Result<(), ApiError> {
    match res.status.class() {
        StatusClass::ClientError => Err(ApiError::ClientError),
        StatusClass::ServerError => Err(ApiError::ServerError),
        _ => Ok(()),
    }
}

// Do the OAUTH stuff for my.mlh.io
// TODO: Something better about the long api urls
fn do_request(code: &str) -> Result<Data, RequestError> {
    let id = env_err!("ID");
    let secret = env_err!("SECRET");
    let redirect = env_err!("REDIRECT");

    let url = format!("https://my.mlh.io/oauth/token?client_id={}&\
                      client_secret={}&code={}&redirect_uri={}&\
                      grant_type=authorization_code",
                      id,
                      secret,
                      code,
                      redirect);

    let client = Client::new();
    let mut res = try!(client.post(&url)
        .header(Connection::close())
        .send());

    try!(check_http_error(&res));

    let mut body = String::new();
    try!(res.read_to_string(&mut body));

    let token: TokenResp = try!(json::decode(&body));
    let url2 = format!("https://my.mlh.io/api/v2/user?access_token={}",
                       token.access_token);

    res = try!(client.get(&url2)
        .header(Connection::close())
        .send());

    try!(check_http_error(&res));

    body = String::new();
    try!(res.read_to_string(&mut body));

    let person_data: Data = try!(json::decode(&body));

    Ok(person_data)
}

// Send a message to slack when a new user signs up
fn slack_send(user: User) -> Result<(), RequestError> {
    let url = env_err!("SLACKURL");
    let client = Client::new();
    let payload = Payload {
        channel: "#signups".to_string(),
        username: "Signup bot".to_string(),
        icon_emoji: ":hackcwru:".to_string(),
        text: format!("{}, a {} major from {}, has signed up!",
                      user.first_name,
                      user.major,
                      user.school.name),
    };
    let payload_str = try!(json::encode(&payload));
    let res = try!(client.post(&url)
        .body(&payload_str)
        .send());

    try!(check_http_error(&res));

    Ok(())
}

// Create the table if it doesn't exist.  Runs on each startup
fn create_table(conn: PooledConnection<PostgresConnectionManager>) {
    let _r =
        conn.execute("CREATE TABLE IF NOT EXISTS person (
                id SERIAL PRIMARY KEY,
                email VARCHAR NOT NULL,
                created_at VARCHAR NOT NULL,
                updated_at VARCHAR NOT NULL,
                first_name VARCHAR NOT NULL,
                last_name VARCHAR NOT NULL,
                graduation VARCHAR,
                major VARCHAR NOT NULL,
                shirt_size VARCHAR NOT NULL,
                dietary_restrictions VARCHAR NOT NULL,
                special_needs VARCHAR,
                date_of_birth VARCHAR NOT NULL,
                gender VARCHAR NOT NULL,
                phone_number VARCHAR NOT NULL,
                school_id integer,
                school_name VARCHAR
                )",
            &[]
        );
    let _r = conn.execute("CREATE TABLE IF NOT EXISTS email (email VARCHAR \
                          PRIMARY KEY)",
                          &[]);
    let _r = conn.execute("ALTER TABLE person ADD COLUMN year integer \
                          NOT NULL DEFAULT 2016",
                          &[]);
    let _r = conn.execute("ALTER TABLE person ALTER COLUMN graduation \
                          DROP NOT NULL",
                          &[]);
    let _r = conn.execute(
            "ALTER TABLE person ADD COLUMN signup_time VARCHAR", &[]
        );
    let _r = conn.execute(
            "ALTER TABLE person ADD COLUMN signed_in
            BOOLEAN NOT NULL DEFAULT FALSE",
        &[]
        );
    let _r = conn.execute(
            "ALTER TABLE person ADD COLUMN signed_in_time VARCHAR", &[]
        );

    let _r = conn.execute(
            "ALTER TABLE person DROP CONSTRAINT person_pkey", &[]
        );

    let _r = conn.execute(
            "ALTER TABLE person ADD PRIMARY KEY (id, year)", &[]
        );
}

fn main() {
    let mut app = Nickel::new();

    let postgres_url = env_err!("DATABASE");
    let dbpool =
        PostgresMiddleware::new(&*postgres_url, TlsMode::None, 5, Box::new(NopErrorHandler))
            .expect("Failed to start PostgresMiddleware");

    create_table(dbpool.pool.clone().get().unwrap());
    app.utilize(dbpool);

    app.get("/callback", middleware! { |request, mut response|
        let conn = request.db_conn();
        let redirect_success = env_err!("FINALREDIRECT");
        let redirect_failure = env_err!("REDIRECTERR");
        let year = env_err!("YEAR").parse::<i32>().unwrap();
        let code = match request.query().get("code") {
            Some(s) => s,
            None => {
                println!("No code, quitting");
                return response.error(StatusCode::BadRequest, "Failed");
            }
        };
        let user_data = match do_request(code) {
            Ok(s) => s,
            Err(e) => {
                println!("Get user data failed with error: {:?}", e);
                return response
                    .error(StatusCode::InternalServerError, "Failed");
            }
        }.data;

        let now = now();
        let now_s = format!("{}", now.rfc3339());
        // TODO: make this a function
        let r = conn.execute(
                "INSERT INTO person (id, email, created_at, updated_at,
                first_name, last_name, major, shirt_size,
                dietary_restrictions, special_needs, date_of_birth, gender,
                phone_number, school_id, school_name, year, signup_time) VALUES
                ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14,
                $15, $16, $17)",
                &[&user_data.id, &user_data.email, &user_data.created_at,
                  &user_data.updated_at, &user_data.first_name,
                  &user_data.last_name, &user_data.major,
                  &user_data.shirt_size, &user_data.dietary_restrictions,
                  &user_data.special_needs, &user_data.date_of_birth,
                  &user_data.gender, &user_data.phone_number,
                  &user_data.school.id, &user_data.school.name, &year, &now_s]
            );
        let redirect = match r {
            Ok(v) => {
                println!("Add to database succeeded with status {:?}", v);
                match slack_send(user_data) {
                    Ok(_) => println!("Slack send worked"),
                    Err(e) => println!("Slack send failed with error: {:?}", e)
                }
                redirect_success
            },
            Err(e) => {
                println!("Add to database failed: {:?}", e);
                redirect_failure
            }
        };

        return response.redirect(&redirect[..]);
    });

    app.get("/start", middleware! { |_req, response|
        let id = env_err!("ID");
        let redirect = env_err!("REDIRECT");
        return response.redirect(
            format!(
                    "http://my.mlh.io/oauth/authorize?client_id={}&\
                    redirect_uri={}&response_type=code",
                    id, redirect
                )
            );
    });

    app.post("/email", middleware! { |req, res|
        let conn = req.db_conn();
        let email = req.param("email").unwrap();
        let redirect_success = env_err!("FINALREDIRECT");
        let redirect_failure = env_err!("REDIRECTERR");
        let r = conn.execute("INSERT INTO email VALEUS $1", &[&email]);
        let redirect = match r {
            Ok(v) => {
                println!("Add to database succeeded with status {:?}", v);
                redirect_success
            },
            Err(e) => {
                println!("Add to database failed: {:?}", e);
                redirect_failure
            }
        };

        return res.redirect(&redirect[..]);
    });

    let bind = match env::var("BIND") {
        Ok(v) => v,
        Err(_) => "127.0.0.1:8080".to_string(),
    };

    let _ = app.listen(&bind[..]);
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeremy Wall <jeremy@marzhillstudios.com>
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
#[macro_use]
extern crate clap;
extern crate ucglib;

use std::cell::RefCell;
use std::error::Error;
use std::fs::File;
use std::io;
use std::path::{Path, PathBuf};
use std::process;
use std::rc::Rc;

use ucglib::build;
use ucglib::build::assets::{Cache, MemoryCache};
use ucglib::build::Val;
use ucglib::convert::traits;
use ucglib::convert::ConverterRegistry;

// TODO(jwall): List the target output types automatically.
fn do_flags<'a, 'b>() -> clap::App<'a, 'b> {
    clap_app!(
        ucg =>
            (version: crate_version!())
            (author: crate_authors!())
            (about: "Universal Configuration Grammar compiler.")
            (@arg nostrict: --("no-strict") "Turn off strict checking.")
            (@subcommand inspect =>
             (about: "Inspect a specific symbol in a ucg file.")
             (@arg sym: --sym +takes_value +required "Specify a specific binding in the ucg file to output.")
             (@arg target: --format +required +takes_value "Inspect output type. (flags, json, env, exec)")
             (@arg INPUT: +required "Input ucg file to inspect symbol from.")
            )
            (@subcommand build =>
             (about: "Build a list of ucg files.")
             (@arg recurse: -r +required conflicts_with[INPUT] "Whether we should recurse in directories or not.")
             (@arg INPUT: ... "Input ucg files or directories to build. If not provided then build the contents of the current directory.")
            )
            (@subcommand test =>
             (about: "Check a list of ucg files for errors and run test assertions.")
             (@arg recurse: -r +required conflicts_with[INPUT] "Whether we should recurse or not.")
             (@arg INPUT: ... "Input ucg files or directories to run test assertions for. If not provided it will scan the current directory for files with _test.ucg")
            )
            (@subcommand converters =>
             (about: "list the available converters")
            )
    )
}

fn run_converter(c: &traits::Converter, v: Rc<Val>, f: Option<&str>) -> traits::Result {
    let mut file: Box<std::io::Write> = match f {
        Some(f) => {
            let mut path_buf = PathBuf::from(f);
            path_buf.set_extension(c.file_ext());
            let new_path = path_buf.to_str().unwrap();
            Box::new(try!(File::create(&new_path)))
        }
        None => Box::new(io::stdout()),
    };
    c.convert(v, file.as_mut())
}

fn build_file(
    file: &str,
    validate: bool,
    strict: bool,
    cache: Rc<RefCell<Cache>>,
) -> Result<build::Builder, Box<Error>> {
    let mut root = PathBuf::from(file);
    if root.is_relative() {
        root = std::env::current_dir().unwrap().join(root);
    }
    let mut builder = build::Builder::new(root.parent().unwrap(), cache);
    builder.set_strict(strict);
    if validate {
        builder.enable_validate_mode();
    }
    try!(builder.build_file(file));
    if validate {
        println!("{}", builder.assert_collector.summary);
    }
    Ok(builder)
}

fn do_validate(file: &str, strict: bool, cache: Rc<RefCell<Cache>>) -> bool {
    println!("Validating {}", file);
    match build_file(file, true, strict, cache) {
        Ok(b) => {
            if b.assert_collector.success {
                println!("File {} Pass\n", file);
            } else {
                println!("File {} Fail\n", file);
                return false;
            }
        }
        Err(msg) => {
            eprintln!("Err: {}", msg);
            return false;
        }
    }
    return true;
}

fn do_compile(
    file: &str,
    strict: bool,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
) -> bool {
    println!("Building {}", file);
    let builder = match build_file(file, false, strict, cache.clone()) {
        Ok(builder) => builder,
        Err(err) => {
            eprintln!("{}", err);
            return false;
        }
    };
    let (typ, val) = match builder.out_lock {
        Some((ref typ, ref val)) => (typ, val.clone()),
        None => {
            eprintln!("Build results in no artifacts.");
            return false;
        }
    };
    match registry.get_converter(typ) {
        Some(converter) => {
            run_converter(converter, val, Some(file)).unwrap();
            eprintln!("Build successful");
            return true;
        }
        None => {
            eprintln!("No such converter {}", typ);
            return false;
        }
    }
}

fn visit_ucg_files(
    path: &Path,
    recurse: bool,
    validate: bool,
    strict: bool,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
) -> Result<bool, Box<Error>> {
    let our_path = String::from(path.to_string_lossy());
    let mut result = true;
    // TODO(jwall): Report the failing files at the bottom.
    let mut summary = String::new();
    if path.is_dir() {
        let mut dir_iter = try!(std::fs::read_dir(path)).peekable();
        loop {
            let entry = match dir_iter.next() {
                Some(e) => e,
                None => {
                    break;
                }
            };
            let next_item = try!(entry);
            let next_path = next_item.path();
            let path_as_string = String::from(next_path.to_string_lossy());
            if next_path.is_dir() && recurse {
                if let Err(e) = visit_ucg_files(
                    &next_path,
                    recurse,
                    validate,
                    strict,
                    cache.clone(),
                    registry,
                ) {
                    eprintln!("{}", e);
                    result = false;
                }
            } else {
                if validate && path_as_string.ends_with("_test.ucg") {
                    if !do_validate(&path_as_string, strict, cache.clone()) {
                        result = false;
                        summary.push_str(format!("{} - FAIL\n", path_as_string).as_str())
                    } else {
                        summary.push_str(format!("{} - PASS\n", path_as_string).as_str())
                    }
                } else if !validate && path_as_string.ends_with(".ucg") {
                    if !do_compile(&path_as_string, strict, cache.clone(), registry) {
                        result = false;
                    }
                }
            }
        }
    } else if validate && our_path.ends_with("_test.ucg") {
        if !do_validate(&our_path, strict, cache) {
            result = false;
            summary.push_str(format!("{} - FAIL\n", our_path).as_str());
        } else {
            summary.push_str(format!("{} - PASS\n", &our_path).as_str());
        }
    } else if !validate {
        if !do_compile(&our_path, strict, cache, registry) {
            result = false;
        }
    }
    if validate && !summary.is_empty() {
        println!("RESULTS:");
        println!("{}", summary);
    }
    Ok(result)
}

fn inspect_command(
    matches: &clap::ArgMatches,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
    strict: bool,
) {
    let file = matches.value_of("INPUT").unwrap();
    let sym = matches.value_of("sym");
    let target = matches.value_of("target").unwrap();
    let root = PathBuf::from(file);
    let mut builder = build::Builder::new(root.parent().unwrap(), cache);
    builder.set_strict(strict);
    match registry.get_converter(target) {
        Some(converter) => {
            // TODO(jwall): We should warn if this is a test file.
            let result = builder.build_file(file);
            if !result.is_ok() {
                eprintln!("{:?}", result.err().unwrap());
                process::exit(1);
            }
            let val = match sym {
                Some(sym_name) => builder.get_out_by_name(sym_name),
                None => builder.last,
            };
            match val {
                Some(value) => {
                    // We use None here because we always output to stdout for an inspect.
                    run_converter(converter, value, None).unwrap();
                    eprintln!("Build successful");
                    process::exit(0);
                }
                None => {
                    eprintln!("Build results in no value.");
                    process::exit(1);
                }
            }
        }
        None => {
            eprintln!("No such converter {}", target);
            process::exit(1);
        }
    }
}

fn build_command(
    matches: &clap::ArgMatches,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
    strict: bool,
) {
    let files = matches.values_of("INPUT");
    let recurse = matches.is_present("recurse");
    let mut ok = true;
    if files.is_none() {
        let curr_dir = std::env::current_dir().unwrap();
        let ok = visit_ucg_files(
            curr_dir.as_path(),
            recurse,
            false,
            strict,
            cache.clone(),
            &registry,
        );
        if let Ok(false) = ok {
            process::exit(1)
        }
        process::exit(0);
    }
    for file in files.unwrap() {
        let pb = PathBuf::from(file);
        if let Ok(false) = visit_ucg_files(&pb, recurse, false, strict, cache.clone(), &registry) {
            ok = false;
        }
    }
    if !ok {
        process::exit(1)
    }
}

fn test_command(
    matches: &clap::ArgMatches,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
    strict: bool,
) {
    let files = matches.values_of("INPUT");
    let recurse = matches.is_present("recurse");
    if files.is_none() {
        let curr_dir = std::env::current_dir().unwrap();
        let ok = visit_ucg_files(
            curr_dir.as_path(),
            recurse,
            true,
            strict,
            cache.clone(),
            &registry,
        );
        if let Ok(false) = ok {
            process::exit(1)
        }
    } else {
        let mut ok = true;
        for file in files.unwrap() {
            let pb = PathBuf::from(file);
            //if pb.is_dir() {
            if let Ok(false) = visit_ucg_files(
                pb.as_path(),
                recurse,
                true,
                strict,
                cache.clone(),
                &registry,
            ) {
                ok = false;
            }
        }
        if !ok {
            process::exit(1)
        }
    }
    process::exit(0);
}

fn converters_command(registry: &ConverterRegistry) {
    println!("Available converters:");
    println!("");
    for (name, c) in registry.get_converter_list().iter() {
        println!("- {}", name);
        println!("  Description: {}", c.description());
        println!("  Output Extension: `.{}`", c.file_ext());
        println!("");
    }
}

fn main() {
    let mut app = do_flags();
    let app_matches = app.clone().get_matches();
    let cache: Rc<RefCell<Cache>> = Rc::new(RefCell::new(MemoryCache::new()));
    let registry = ConverterRegistry::make_registry();
    let strict = if app_matches.is_present("nostrict") {
        false
    } else {
        true
    };
    if let Some(matches) = app_matches.subcommand_matches("inspect") {
        inspect_command(matches, cache, &registry, strict);
    } else if let Some(matches) = app_matches.subcommand_matches("build") {
        build_command(matches, cache, &registry, strict);
    } else if let Some(matches) = app_matches.subcommand_matches("test") {
        test_command(matches, cache, &registry, strict);
    } else if let Some(_) = app_matches.subcommand_matches("converters") {
        converters_command(&registry)
    } else {
        app.print_help().unwrap();
        println!("");
    }
}
<commit_msg>Recurse flag can be used with flag input as arguments.<commit_after>// Copyright 2017 Jeremy Wall <jeremy@marzhillstudios.com>
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
#[macro_use]
extern crate clap;
extern crate ucglib;

use std::cell::RefCell;
use std::error::Error;
use std::fs::File;
use std::io;
use std::path::{Path, PathBuf};
use std::process;
use std::rc::Rc;

use ucglib::build;
use ucglib::build::assets::{Cache, MemoryCache};
use ucglib::build::Val;
use ucglib::convert::traits;
use ucglib::convert::ConverterRegistry;

// TODO(jwall): List the target output types automatically.
fn do_flags<'a, 'b>() -> clap::App<'a, 'b> {
    clap_app!(
        ucg =>
            (version: crate_version!())
            (author: crate_authors!())
            (about: "Universal Configuration Grammar compiler.")
            (@arg nostrict: --("no-strict") "Turn off strict checking.")
            (@subcommand inspect =>
             (about: "Inspect a specific symbol in a ucg file.")
             (@arg sym: --sym +takes_value +required "Specify a specific binding in the ucg file to output.")
             (@arg target: --format +required +takes_value "Inspect output type. (flags, json, env, exec)")
             (@arg INPUT: +required "Input ucg file to inspect symbol from.")
            )
            (@subcommand build =>
             (about: "Build a list of ucg files.")
             (@arg recurse: -r +required "Whether we should recurse in directories or not.")
             (@arg INPUT: ... "Input ucg files or directories to build. If not provided then build the contents of the current directory.")
            )
            (@subcommand test =>
             (about: "Check a list of ucg files for errors and run test assertions.")
             (@arg recurse: -r +required "Whether we should recurse or not.")
             (@arg INPUT: ... "Input ucg files or directories to run test assertions for. If not provided it will scan the current directory for files with _test.ucg")
            )
            (@subcommand converters =>
             (about: "list the available converters")
            )
    )
}

fn run_converter(c: &traits::Converter, v: Rc<Val>, f: Option<&str>) -> traits::Result {
    let mut file: Box<std::io::Write> = match f {
        Some(f) => {
            let mut path_buf = PathBuf::from(f);
            path_buf.set_extension(c.file_ext());
            let new_path = path_buf.to_str().unwrap();
            Box::new(try!(File::create(&new_path)))
        }
        None => Box::new(io::stdout()),
    };
    c.convert(v, file.as_mut())
}

fn build_file(
    file: &str,
    validate: bool,
    strict: bool,
    cache: Rc<RefCell<Cache>>,
) -> Result<build::Builder, Box<Error>> {
    let mut root = PathBuf::from(file);
    if root.is_relative() {
        root = std::env::current_dir().unwrap().join(root);
    }
    let mut builder = build::Builder::new(root.parent().unwrap(), cache);
    builder.set_strict(strict);
    if validate {
        builder.enable_validate_mode();
    }
    try!(builder.build_file(file));
    if validate {
        println!("{}", builder.assert_collector.summary);
    }
    Ok(builder)
}

fn do_validate(file: &str, strict: bool, cache: Rc<RefCell<Cache>>) -> bool {
    println!("Validating {}", file);
    match build_file(file, true, strict, cache) {
        Ok(b) => {
            if b.assert_collector.success {
                println!("File {} Pass\n", file);
            } else {
                println!("File {} Fail\n", file);
                return false;
            }
        }
        Err(msg) => {
            eprintln!("Err: {}", msg);
            return false;
        }
    }
    return true;
}

fn do_compile(
    file: &str,
    strict: bool,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
) -> bool {
    println!("Building {}", file);
    let builder = match build_file(file, false, strict, cache.clone()) {
        Ok(builder) => builder,
        Err(err) => {
            eprintln!("{}", err);
            return false;
        }
    };
    let (typ, val) = match builder.out_lock {
        Some((ref typ, ref val)) => (typ, val.clone()),
        None => {
            eprintln!("Build results in no artifacts.");
            return false;
        }
    };
    match registry.get_converter(typ) {
        Some(converter) => {
            run_converter(converter, val, Some(file)).unwrap();
            eprintln!("Build successful");
            return true;
        }
        None => {
            eprintln!("No such converter {}", typ);
            return false;
        }
    }
}

fn visit_ucg_files(
    path: &Path,
    recurse: bool,
    validate: bool,
    strict: bool,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
) -> Result<bool, Box<Error>> {
    let our_path = String::from(path.to_string_lossy());
    let mut result = true;
    // TODO(jwall): Report the failing files at the bottom.
    let mut summary = String::new();
    if path.is_dir() {
        let mut dir_iter = try!(std::fs::read_dir(path)).peekable();
        loop {
            let entry = match dir_iter.next() {
                Some(e) => e,
                None => {
                    break;
                }
            };
            let next_item = try!(entry);
            let next_path = next_item.path();
            let path_as_string = String::from(next_path.to_string_lossy());
            if next_path.is_dir() && recurse {
                if let Err(e) = visit_ucg_files(
                    &next_path,
                    recurse,
                    validate,
                    strict,
                    cache.clone(),
                    registry,
                ) {
                    eprintln!("{}", e);
                    result = false;
                }
            } else {
                if validate && path_as_string.ends_with("_test.ucg") {
                    if !do_validate(&path_as_string, strict, cache.clone()) {
                        result = false;
                        summary.push_str(format!("{} - FAIL\n", path_as_string).as_str())
                    } else {
                        summary.push_str(format!("{} - PASS\n", path_as_string).as_str())
                    }
                } else if !validate && path_as_string.ends_with(".ucg") {
                    if !do_compile(&path_as_string, strict, cache.clone(), registry) {
                        result = false;
                    }
                }
            }
        }
    } else if validate && our_path.ends_with("_test.ucg") {
        if !do_validate(&our_path, strict, cache) {
            result = false;
            summary.push_str(format!("{} - FAIL\n", our_path).as_str());
        } else {
            summary.push_str(format!("{} - PASS\n", &our_path).as_str());
        }
    } else if !validate {
        if !do_compile(&our_path, strict, cache, registry) {
            result = false;
        }
    }
    if validate && !summary.is_empty() {
        println!("RESULTS:");
        println!("{}", summary);
    }
    Ok(result)
}

fn inspect_command(
    matches: &clap::ArgMatches,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
    strict: bool,
) {
    let file = matches.value_of("INPUT").unwrap();
    let sym = matches.value_of("sym");
    let target = matches.value_of("target").unwrap();
    let root = PathBuf::from(file);
    let mut builder = build::Builder::new(root.parent().unwrap(), cache);
    builder.set_strict(strict);
    match registry.get_converter(target) {
        Some(converter) => {
            // TODO(jwall): We should warn if this is a test file.
            let result = builder.build_file(file);
            if !result.is_ok() {
                eprintln!("{:?}", result.err().unwrap());
                process::exit(1);
            }
            let val = match sym {
                Some(sym_name) => builder.get_out_by_name(sym_name),
                None => builder.last,
            };
            match val {
                Some(value) => {
                    // We use None here because we always output to stdout for an inspect.
                    run_converter(converter, value, None).unwrap();
                    eprintln!("Build successful");
                    process::exit(0);
                }
                None => {
                    eprintln!("Build results in no value.");
                    process::exit(1);
                }
            }
        }
        None => {
            eprintln!("No such converter {}", target);
            process::exit(1);
        }
    }
}

fn build_command(
    matches: &clap::ArgMatches,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
    strict: bool,
) {
    let files = matches.values_of("INPUT");
    let recurse = matches.is_present("recurse");
    let mut ok = true;
    if files.is_none() {
        let curr_dir = std::env::current_dir().unwrap();
        let ok = visit_ucg_files(
            curr_dir.as_path(),
            recurse,
            false,
            strict,
            cache.clone(),
            &registry,
        );
        if let Ok(false) = ok {
            process::exit(1)
        }
        process::exit(0);
    }
    for file in files.unwrap() {
        let pb = PathBuf::from(file);
        if let Ok(false) = visit_ucg_files(&pb, recurse, false, strict, cache.clone(), &registry) {
            ok = false;
        }
    }
    if !ok {
        process::exit(1)
    }
}

fn test_command(
    matches: &clap::ArgMatches,
    cache: Rc<RefCell<Cache>>,
    registry: &ConverterRegistry,
    strict: bool,
) {
    let files = matches.values_of("INPUT");
    let recurse = matches.is_present("recurse");
    if files.is_none() {
        let curr_dir = std::env::current_dir().unwrap();
        let ok = visit_ucg_files(
            curr_dir.as_path(),
            recurse,
            true,
            strict,
            cache.clone(),
            &registry,
        );
        if let Ok(false) = ok {
            process::exit(1)
        }
    } else {
        let mut ok = true;
        for file in files.unwrap() {
            let pb = PathBuf::from(file);
            //if pb.is_dir() {
            if let Ok(false) = visit_ucg_files(
                pb.as_path(),
                recurse,
                true,
                strict,
                cache.clone(),
                &registry,
            ) {
                ok = false;
            }
        }
        if !ok {
            process::exit(1)
        }
    }
    process::exit(0);
}

fn converters_command(registry: &ConverterRegistry) {
    println!("Available converters:");
    println!("");
    for (name, c) in registry.get_converter_list().iter() {
        println!("- {}", name);
        println!("  Description: {}", c.description());
        println!("  Output Extension: `.{}`", c.file_ext());
        println!("");
    }
}

fn main() {
    let mut app = do_flags();
    let app_matches = app.clone().get_matches();
    let cache: Rc<RefCell<Cache>> = Rc::new(RefCell::new(MemoryCache::new()));
    let registry = ConverterRegistry::make_registry();
    let strict = if app_matches.is_present("nostrict") {
        false
    } else {
        true
    };
    if let Some(matches) = app_matches.subcommand_matches("inspect") {
        inspect_command(matches, cache, &registry, strict);
    } else if let Some(matches) = app_matches.subcommand_matches("build") {
        build_command(matches, cache, &registry, strict);
    } else if let Some(matches) = app_matches.subcommand_matches("test") {
        test_command(matches, cache, &registry, strict);
    } else if let Some(_) = app_matches.subcommand_matches("converters") {
        converters_command(&registry)
    } else {
        app.print_help().unwrap();
        println!("");
    }
}
<|endoftext|>


<commit_before>extern crate diff;
extern crate regex;

use std::fs::File;
use std::process::{Command, Stdio};
use std::collections::BinaryHeap;
use std::fs;
use regex::Regex;
use std::io::{self, Read, Write};

#[derive(Debug)]
struct Testcase {
    input: String,
    output: String,
}

enum Error {
    Io(io::Error),
    TooManyInputs((BinaryHeap<String>, BinaryHeap<String>)),
    TooManyOutputs((BinaryHeap<String>, BinaryHeap<String>)),
}

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

fn get_testcases() -> Result<Vec<Testcase>, Error> {
    let input_re = Regex::new(r"input.*").unwrap();
    let output_re = Regex::new(r"output.*").unwrap();

    let mut testcases = vec![];
    let mut inputs = BinaryHeap::new();
    let mut outputs = BinaryHeap::new();

    for entry in fs::read_dir(".")? {
        let path = entry?.path().into_os_string();
        if let Some(input) = input_re.captures(path.to_str().unwrap()) {
            inputs.push(input[0].to_string());
        }
        if let Some(output) = output_re.captures(path.to_str().unwrap()) {
            outputs.push(output[0].to_string());
        }
    }
    if inputs.len() < outputs.len() {
        return Err(Error::TooManyOutputs((inputs, outputs)));
    } else if inputs.len() > outputs.len() {
        return Err(Error::TooManyInputs((inputs, outputs)));
    }
    while !inputs.is_empty() {
        testcases.push(Testcase {
            input: inputs.pop().unwrap(),
            output: outputs.pop().unwrap(),
        });
    }
    Ok(testcases)
}

fn run_testcase(binary: &str, testcase: &Testcase) -> Result<(), io::Error> {
    let process = Command::new(binary).stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()?;
    let mut f = File::open(&testcase.input)?;
    let mut inp = String::new();
    f.read_to_string(&mut inp)?;
    process.stdin.unwrap().write_all(inp.as_bytes())?;

    let mut outf = File::open(&testcase.output)?;
    let mut out_expected = String::new();
    outf.read_to_string(&mut out_expected)?;

    let mut out = String::new();
    process.stdout.unwrap().read_to_string(&mut out)?;

    if out == out_expected {
        Ok(())
    } else {
        let mut s = String::new();
        for diff in diff::lines(&out, &out_expected) {
            s += &match diff {
                diff::Result::Left(l)    => format!("-{}\n", l),
                diff::Result::Both(l, _) => format!(" {}\n", l),
                diff::Result::Right(r)   => format!("+{}\n", r),
            };
        }
        Err(io::Error::new(io::ErrorKind::Other, s))
    }
}

fn main() {
    let binary = if let Some(name) = std::env::args().nth(2) {
        name
    } else {
        String::from("./a.out")
    };

    if !std::path::Path::new(&binary).exists() {
        println!("Binary {} does not exist", binary);
        std::process::exit(1);
    }

    let mut failed_testcases = vec![];
    match get_testcases() {
        Ok(testcases) => {
            println!("{} testcases found", testcases.len());
            for testcase in testcases {
                match run_testcase(&binary, &testcase) {
                    Ok(_) => print!("."),
                    Err(e) => {
                        failed_testcases.push((testcase, e));
                        print!("E");
                    }
                }
            }
            println!();
        }
        Err(Error::Io(e)) => println!("I/O error: {}", e),
        Err(Error::TooManyInputs((i, o))) => {
            println!("There are more inputs than outputs.
Inputs : {:?}
Outputs: {:?}",
                     i,
                     o);
        }
        Err(Error::TooManyOutputs((i, o))) => {
            println!("There are more outputs than inputs.
Inputs : {:?}
Outputs: {:?}",
                     i,
                     o);
        }
    }
    if !failed_testcases.is_empty() {
        println!("Failed testcases:");
        for (testcase, e) in failed_testcases {
            println!("{:?}:\n{}\n", testcase, e);
        }
        std::process::exit(1);
    } else {
        println!("OK");
    }
}
<commit_msg>Fix argument number<commit_after>extern crate diff;
extern crate regex;

use std::fs::File;
use std::process::{Command, Stdio};
use std::collections::BinaryHeap;
use std::fs;
use regex::Regex;
use std::io::{self, Read, Write};

#[derive(Debug)]
struct Testcase {
    input: String,
    output: String,
}

enum Error {
    Io(io::Error),
    TooManyInputs((BinaryHeap<String>, BinaryHeap<String>)),
    TooManyOutputs((BinaryHeap<String>, BinaryHeap<String>)),
}

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

fn get_testcases() -> Result<Vec<Testcase>, Error> {
    let input_re = Regex::new(r"input.*").unwrap();
    let output_re = Regex::new(r"output.*").unwrap();

    let mut testcases = vec![];
    let mut inputs = BinaryHeap::new();
    let mut outputs = BinaryHeap::new();

    for entry in fs::read_dir(".")? {
        let path = entry?.path().into_os_string();
        if let Some(input) = input_re.captures(path.to_str().unwrap()) {
            inputs.push(input[0].to_string());
        }
        if let Some(output) = output_re.captures(path.to_str().unwrap()) {
            outputs.push(output[0].to_string());
        }
    }
    if inputs.len() < outputs.len() {
        return Err(Error::TooManyOutputs((inputs, outputs)));
    } else if inputs.len() > outputs.len() {
        return Err(Error::TooManyInputs((inputs, outputs)));
    }
    while !inputs.is_empty() {
        testcases.push(Testcase {
            input: inputs.pop().unwrap(),
            output: outputs.pop().unwrap(),
        });
    }
    Ok(testcases)
}

fn run_testcase(binary: &str, testcase: &Testcase) -> Result<(), io::Error> {
    let process = Command::new(binary).stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()?;
    let mut f = File::open(&testcase.input)?;
    let mut inp = String::new();
    f.read_to_string(&mut inp)?;
    process.stdin.unwrap().write_all(inp.as_bytes())?;

    let mut outf = File::open(&testcase.output)?;
    let mut out_expected = String::new();
    outf.read_to_string(&mut out_expected)?;

    let mut out = String::new();
    process.stdout.unwrap().read_to_string(&mut out)?;

    if out == out_expected {
        Ok(())
    } else {
        let mut s = String::new();
        for diff in diff::lines(&out, &out_expected) {
            s += &match diff {
                diff::Result::Left(l)    => format!("-{}\n", l),
                diff::Result::Both(l, _) => format!(" {}\n", l),
                diff::Result::Right(r)   => format!("+{}\n", r),
            };
        }
        Err(io::Error::new(io::ErrorKind::Other, s))
    }
}

fn main() {
    let binary = if let Some(name) = std::env::args().nth(1) {
        name
    } else {
        String::from("./a.out")
    };

    if !std::path::Path::new(&binary).exists() {
        println!("Binary {} does not exist", binary);
        std::process::exit(1);
    }

    let mut failed_testcases = vec![];
    match get_testcases() {
        Ok(testcases) => {
            println!("{} testcases found", testcases.len());
            for testcase in testcases {
                match run_testcase(&binary, &testcase) {
                    Ok(_) => print!("."),
                    Err(e) => {
                        failed_testcases.push((testcase, e));
                        print!("E");
                    }
                }
            }
            println!();
        }
        Err(Error::Io(e)) => println!("I/O error: {}", e),
        Err(Error::TooManyInputs((i, o))) => {
            println!("There are more inputs than outputs.
Inputs : {:?}
Outputs: {:?}",
                     i,
                     o);
        }
        Err(Error::TooManyOutputs((i, o))) => {
            println!("There are more outputs than inputs.
Inputs : {:?}
Outputs: {:?}",
                     i,
                     o);
        }
    }
    if !failed_testcases.is_empty() {
        println!("Failed testcases:");
        for (testcase, e) in failed_testcases {
            println!("{:?}:\n{}\n", testcase, e);
        }
        std::process::exit(1);
    } else {
        println!("OK");
    }
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate numbat;

mod forza;
mod plugins;

use std::thread;
use std::collections::BTreeMap;
use forza::ForzaPlugin;

fn start_plugin<T: ForzaPlugin>(mut plugin: T) {
  plugin.start();
}

fn main() {
  let mut emitter = numbat::Emitter::new(BTreeMap::new(), "host");
  emitter.connect(&match std::env::var_os("METRICS") {
    Some(url) => url.into_string().expect("expected METRICS to be valid UTF-8"),
    None => String::from("tcp://127.0.0.1:1337")
  });

  let heartbeat = plugins::heartbeat::Heartbeat::new(emitter.clone());
  thread::spawn(|| {
    start_plugin(heartbeat);
  });

  let memory = plugins::memory::Memory::new(emitter.clone());
  thread::spawn(|| {
    start_plugin(memory);
  });

  let load_average = plugins::load_average::LoadAverage::new(emitter.clone());
  thread::spawn(|| {
    start_plugin(load_average);
  }).join();
}
<commit_msg>Emit a metric on Forza start<commit_after>extern crate libc;
extern crate numbat;

mod forza;
mod plugins;

use std::thread;
use std::collections::BTreeMap;
use forza::ForzaPlugin;

fn start_plugin<T: ForzaPlugin>(mut plugin: T) {
  plugin.start();
}

fn main() {
  let mut emitter = numbat::Emitter::new(BTreeMap::new(), "host");
  emitter.connect(&match std::env::var_os("METRICS") {
    Some(url) => url.into_string().expect("expected METRICS to be valid UTF-8"),
    None => String::from("tcp://127.0.0.1:1337")
  });

  emitter.emit_name("forza.start");

  let heartbeat = plugins::heartbeat::Heartbeat::new(emitter.clone());
  thread::spawn(|| {
    start_plugin(heartbeat);
  });

  let memory = plugins::memory::Memory::new(emitter.clone());
  thread::spawn(|| {
    start_plugin(memory);
  });

  let load_average = plugins::load_average::LoadAverage::new(emitter.clone());
  thread::spawn(|| {
    start_plugin(load_average);
  }).join();
}
<|endoftext|>


<commit_before>extern crate webextension_rust_template as protocol;
use std::io;

fn main() {
    loop {
        let f = protocol::Input::Stdin(io::stdin());
        let message = protocol::read(f);
        let output = protocol::Output::Stdout(io::stdout());
        protocol::write(output, message.to_string());
    }
}
<commit_msg>separate thread reads from stdin<commit_after>extern crate webextension_rust_template as protocol;
use std::io;
use std::thread;
use std::sync::mpsc;

fn main() {
    let (tx, rx) = mpsc::channel();
    let sender = tx.clone();

    // thread reads from stdin and send messages to main thread
    thread::spawn(move || {
        loop {
            let f = protocol::Input::Stdin(io::stdin());
            let message = protocol::read(f);
            sender.send(message).unwrap();
        }
    });

    loop {
        let message = rx.recv().unwrap();
        let output = protocol::Output::Stdout(io::stdout());
        protocol::write(output, message.to_string());
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use std::env;
use std::net::TcpListener;
use std::os::unix::io::{AsRawFd,FromRawFd};
use std::process::{Command,Stdio};

fn main() {
    // process arguments
    let mut args = env::args();
    args.next().unwrap();
    let command = args.next().unwrap();
    let command_args = args.collect::<Vec<_>>();

    // open tcp socket
    let listener = TcpListener::bind("[::1]:0").unwrap();
    let port = listener.local_addr().unwrap().port();

    // convert port to string for later use
    // dash prefix forces telnet protocol
    let str_port = format!("-{}", port);

    // spawn telnet
    let mut telnet = Command::new("telnet")
    .args(&["::1", "--", &str_port])
    .spawn().unwrap();

    // wait for telnet to connect
    let (stream, _) = listener.accept().unwrap();
    drop(listener);
    let stdin = unsafe { Stdio::from_raw_fd(libc::dup(stream.as_raw_fd())) };
    let stdout = unsafe { Stdio::from_raw_fd(libc::dup(stream.as_raw_fd())) };
    let stderr = unsafe { Stdio::from_raw_fd(libc::dup(stream.as_raw_fd())) };

    // spawn application
    let mut application = Command::new(command)
    .args(&command_args)
    .stdin(stdin)
    .stdout(stdout)
    .stderr(stderr)
    .spawn().unwrap();

    // wait for all processes to terminate
    telnet.wait().unwrap();
    application.wait().unwrap();
}
<commit_msg>isolate unsafe code<commit_after>extern crate libc;

use std::env;
use std::net::TcpListener;
use std::os::unix::io::{AsRawFd,FromRawFd};
use std::process::{Command,Stdio};

fn stdio_dup_from<T: AsRawFd>(x: &T) -> Stdio {
    unsafe { Stdio::from_raw_fd(libc::dup(x.as_raw_fd())) }
}

fn main() {
    // process arguments
    let mut args = env::args();
    args.next().unwrap();
    let command = args.next().unwrap();
    let command_args = args.collect::<Vec<_>>();

    // open tcp socket
    let listener = TcpListener::bind("[::1]:0").unwrap();
    let port = listener.local_addr().unwrap().port();

    // convert port to string for later use
    // dash prefix forces telnet protocol
    let str_port = format!("-{}", port);

    // spawn telnet
    let mut telnet = Command::new("telnet")
    .args(&["::1", "--", &str_port])
    .spawn().unwrap();

    // wait for telnet to connect
    let (stream, _) = listener.accept().unwrap();
    drop(listener);
    let stdin = stdio_dup_from(&stream);
    let stdout = stdio_dup_from(&stream);
    let stderr = stdio_dup_from(&stream);

    // spawn application
    let mut application = Command::new(command)
    .args(&command_args)
    .stdin(stdin)
    .stdout(stdout)
    .stderr(stderr)
    .spawn().unwrap();

    // wait for all processes to terminate
    telnet.wait().unwrap();
    application.wait().unwrap();
}
<|endoftext|>