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

#![allow(clippy::redundant_field_names)]

use clap::{Arg, SubCommand};
use evebox::logger::{self, log};
use evebox::version;

#[tokio::main]
async fn main() {
    if let Err(err) = _main().await {
        log::error!("{}", err);
        std::process::exit(1);
    }
}

async fn _main() -> Result<(), Box<dyn std::error::Error>> {
    let parser = clap::App::new("EveBox")
        .version(std::env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::with_name("verbose")
                .long("verbose")
                .short("v")
                .multiple(true)
                .global(true)
                .help("Increase verbosity"),
        )
        .arg(
            Arg::with_name("data-directory")
                .long("data-directory")
                .short("D")
                .takes_value(true)
                .value_name("DIR")
                .help("Data directory")
                .global(true),
        )
        .subcommand(clap::SubCommand::with_name("version").about("Display version"));

    let sqlite_import = SubCommand::with_name("sqlite-import")
        .about("Import to SQLite")
        .arg(
            Arg::with_name("oneshot")
                .long("oneshot")
                .help("One shot mode (exit on EOF"),
        )
        .arg(
            Arg::with_name("end")
                .long("end")
                .help("Start at end of file"),
        )
        .arg(Arg::with_name("INPUT").required(true).index(1));

    let server = clap::SubCommand::with_name("server")
        .about("EveBox Server")
        .arg(
            clap::Arg::with_name("config")
                .long("config")
                .short("c")
                .takes_value(true)
                .value_name("FILE")
                .help("Configuration filename"),
        )
        .arg(
            clap::Arg::with_name("http.host")
                .long("host")
                .value_name("HOSTNAME")
                .takes_value(true)
                .default_value("127.0.0.1")
                .help("Hostname/IP address to bind to"),
        )
        .arg(
            clap::Arg::with_name("http.port")
                .short("p")
                .long("port")
                .value_name("PORT")
                .takes_value(true)
                .default_value("5636")
                .help("Port to bind to"),
        )
        .arg(
            clap::Arg::with_name("database.elasticsearch.url")
                .short("e")
                .long("elasticsearch")
                .takes_value(true)
                .value_name("URL")
                .default_value("http://localhost:9200")
                .help("Elastic Search URL"),
        )
        .arg(
            clap::Arg::with_name("database.elasticsearch.index")
                .short("i")
                .long("index")
                .takes_value(true)
                .default_value("logstash")
                .value_name("INDEX")
                .help("Elastic Search index prefix"),
        )
        .arg(
            Arg::with_name("database.elasticsearch.no-index-suffix")
                .long("no-index-suffix")
                .takes_value(false)
                .help("Do not add a suffix to the index name"),
        )
        .arg(
            Arg::with_name("database.elasticsearch.ecs")
                .long("ecs")
                .help("Enable Elastic ECS support"),
        )
        .arg(
            Arg::with_name("database.type")
                .long("database")
                .aliases(&["datastore"])
                .takes_value(true)
                .value_name("DATABASE")
                .default_value("elasticsearch")
                .help("Database type"),
        )
        .arg(
            Arg::with_name("sqlite-filename")
                .long("sqlite-filename")
                .takes_value(true)
                .value_name("FILENAME")
                .default_value("events.sqlite")
                .help("SQLite events filename"),
        )
        .arg(
            clap::Arg::with_name("no-check-certificate")
                .short("k")
                .long("no-check-certificate")
                .help("Disable TLS certificate validation"),
        )
        .arg(
            Arg::with_name("http.request-logging")
                .long("http-request-logging")
                .env("EVEBOX_HTTP_REQUEST_LOGGING")
                .help("Enable HTTP access logging"),
        )
        .arg(
            Arg::with_name("http.tls.enabled")
                .long("tls")
                .help("Enable TLS"),
        )
        .arg(
            Arg::with_name("http.tls.certificate")
                .long("tls-cert")
                .takes_value(true)
                .value_name("FILENAME")
                .help("TLS certificate filename"),
        )
        .arg(
            Arg::with_name("http.tls.key")
                .long("tls-key")
                .takes_value(true)
                .value_name("FILENAME")
                .help("TLS key filename"),
        )
        .arg(
            Arg::with_name("input.filename")
                .long("input")
                .takes_value(true)
                .value_name("FILENAME")
                .help("Input Eve file to read"),
        )
        .arg(
            Arg::with_name("end")
                .long("end")
                .help("Read from end (tail) of input file"),
        )
        .arg(
            Arg::with_name("input-start")
                .long("input-start")
                .hidden(true),
        );

    let agent = SubCommand::with_name("agent")
        .about("EveBox Agent")
        .arg(
            Arg::with_name("config")
                .short("c")
                .long("config")
                .takes_value(true)
                .help("Configuration file"),
        )
        .arg(
            Arg::with_name("server.url")
                .long("server")
                .takes_value(true)
                .value_name("URL")
                .help("EveBox Server URL"),
        )
        .arg(Arg::with_name("geoip.enabled").long("enable-geoip"))
        .arg(
            Arg::with_name("stdout")
                .long("stdout")
                .help("Print events to stdout"),
        )
        .arg(
            Arg::with_name("bookmark-directory")
                .long("bookmark-directory")
                .takes_value(true)
                .hidden(true),
        );

    let oneshot = SubCommand::with_name("oneshot")
        .about("Import a single eve.json and review in EveBox")
        .arg(
            Arg::with_name("limit")
                .long("limit")
                .takes_value(true)
                .help("Limit the number of events read"),
        )
        .arg(
            Arg::with_name("no-open")
                .long("no-open")
                .help("Don't open browser"),
        )
        .arg(
            Arg::with_name("no-wait")
                .long("no-wait")
                .help("Don't wait for events to load"),
        )
        .arg(
            Arg::with_name("database-filename")
                .long("database-filename")
                .takes_value(true)
                .default_value("./oneshot.sqlite")
                .help("Database filename"),
        )
        // --host, but keep th name as http.host to be campatible with the
        // EVEBOX_HTTP_HOST environment variable.
        .arg(
            clap::Arg::with_name("http.host")
                .long("host")
                .value_name("HOSTNAME")
                .takes_value(true)
                .default_value("127.0.0.1")
                .help("Hostname/IP address to bind to"),
        )
        .arg(Arg::with_name("INPUT").required(true).index(1));

    let elastic_import = SubCommand::with_name("elastic-import")
        .alias("esimport")
        .about("Import to Elastic Search")
        .arg(
            Arg::with_name("config")
                .short("c")
                .long("config")
                .takes_value(true)
                .help("Configuration file"),
        )
        .arg(
            Arg::with_name("oneshot")
                .long("oneshot")
                .help("One shot mode (exit on EOF)"),
        )
        .arg(
            Arg::with_name("end")
                .long("end")
                .help("Start at end of file"),
        )
        .arg(
            clap::Arg::with_name("elasticsearch")
                .short("e")
                .long("elasticsearch")
                .takes_value(true)
                .default_value("http://localhost:9200")
                .help("Elastic Search URL"),
        )
        .arg(
            Arg::with_name("index")
                .long("index")
                .takes_value(true)
                .default_value("logstash")
                .help("Elastic Search index prefix"),
        )
        .arg(
            Arg::with_name("no-index-suffix")
                .long("no-index-suffix")
                .takes_value(false)
                .help("Do not add a suffix to the index name"),
        )
        .arg(
            Arg::with_name("bookmark")
                .long("bookmark")
                .help("Enable bookmarking"),
        )
        .arg(
            Arg::with_name("bookmark-filename")
                .long("bookmark-filename")
                .takes_value(true)
                //.default_value(".bookmark")
                .default_value("")
                .help("Bookmark filename"),
        )
        .arg(
            Arg::with_name("bookmark-dir")
                .long("bookmark-dir")
                .takes_value(true)
                .default_value(".")
                .help("Bookmark directory"),
        )
        .arg(
            Arg::with_name("stdout")
                .long("stdout")
                .help("Print events to stdout"),
        )
        .arg(
            Arg::with_name("username")
                .long("username")
                .short("u")
                .takes_value(true)
                .help("Elasticsearch username"),
        )
        .arg(
            Arg::with_name("password")
                .long("password")
                .short("p")
                .takes_value(true)
                .help("Elasticsearch password"),
        )
        .arg(
            clap::Arg::with_name(evebox::commands::elastic_import::NO_CHECK_CERTIFICATE)
                .short("k")
                .long("no-check-certificate")
                .help("Disable TLS certificate validation"),
        )
        .arg(
            Arg::with_name("geoip.disable")
                .long("no-geoip")
                .help("Disable GeoIP"),
        )
        .arg(
            Arg::with_name("geoip.database-filename")
                .long("geoip-database")
                .takes_value(true)
                .value_name("filename")
                .help("GeoIP database filename"),
        )
        .arg(
            Arg::with_name("input")
                .multiple(true)
                .value_name("INPUT")
                .index(1),
        );

    let mut parser = parser
        .subcommand(server)
        .subcommand(elastic_import)
        .subcommand(oneshot)
        .subcommand(agent)
        .subcommand(sqlite_import)
        .subcommand(evebox::commands::config::config_subcommand())
        .subcommand(elastic_debug());

    let matches = parser.clone().get_matches();

    // Initialize logging.
    let log_level = {
        let verbosity = matches.occurrences_of("verbose");
        if verbosity > 1 {
            tracing::Level::TRACE
        } else if verbosity > 0 {
            tracing::Level::DEBUG
        } else {
            tracing::Level::INFO
        }
    };
    logger::init_logger(log_level);
    logger::init_stdlog();

    let rc: anyhow::Result<()> = match matches.subcommand() {
        ("server", Some(args)) => {
            evebox::server::main(args).await?;
            Ok(())
        }
        ("version", _) => {
            version::print_version();
            Ok(())
        }
        ("sqlite-import", Some(args)) => evebox::commands::sqlite_import::main(args).await,
        ("elastic-import", Some(args)) => {
            if let Err(err) = evebox::commands::elastic_import::main(&args).await {
                log::error!("{}", err);
                std::process::exit(1);
            }
            Ok(())
        }
        ("elastic-debug", Some(args)) => evebox::commands::elastic_debug::main(args).await,
        ("oneshot", Some(args)) => evebox::commands::oneshot::main(&args).await,
        ("agent", Some(args)) => evebox::agent::main(&args).await,
        ("config", Some(args)) => evebox::commands::config::main(&args),
        ("", None) => {
            parser.print_help().ok();
            // print_help doesn't output a new line at the end, so fix that up...
            println!();
            std::process::exit(1);
        }
        (command, _) => {
            log::error!("command \"{}\" not implemented yet", command);
            std::process::exit(1);
        }
    };
    if let Err(err) = rc {
        log::error!("error: {}", err);
        std::process::exit(1);
    } else {
        Ok(())
    }
}

fn elastic_debug() -> clap::App<'static, 'static> {
    SubCommand::with_name("elastic-debug").arg(
        Arg::with_name("elasticsearch")
            .long("elasticsearch")
            .short("e")
            .takes_value(true)
            .default_value("127.0.0.1"),
    )
}
<commit_msg>elastic-import: fix --no-geoip command line<commit_after>// Copyright (C) 2020 Jason Ish
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#![allow(clippy::redundant_field_names)]

use clap::{Arg, SubCommand};
use evebox::logger::{self, log};
use evebox::version;

#[tokio::main]
async fn main() {
    if let Err(err) = _main().await {
        log::error!("{}", err);
        std::process::exit(1);
    }
}

async fn _main() -> Result<(), Box<dyn std::error::Error>> {
    let parser = clap::App::new("EveBox")
        .version(std::env!("CARGO_PKG_VERSION"))
        .arg(
            Arg::with_name("verbose")
                .long("verbose")
                .short("v")
                .multiple(true)
                .global(true)
                .help("Increase verbosity"),
        )
        .arg(
            Arg::with_name("data-directory")
                .long("data-directory")
                .short("D")
                .takes_value(true)
                .value_name("DIR")
                .help("Data directory")
                .global(true),
        )
        .subcommand(clap::SubCommand::with_name("version").about("Display version"));

    let sqlite_import = SubCommand::with_name("sqlite-import")
        .about("Import to SQLite")
        .arg(
            Arg::with_name("oneshot")
                .long("oneshot")
                .help("One shot mode (exit on EOF"),
        )
        .arg(
            Arg::with_name("end")
                .long("end")
                .help("Start at end of file"),
        )
        .arg(Arg::with_name("INPUT").required(true).index(1));

    let server = clap::SubCommand::with_name("server")
        .about("EveBox Server")
        .arg(
            clap::Arg::with_name("config")
                .long("config")
                .short("c")
                .takes_value(true)
                .value_name("FILE")
                .help("Configuration filename"),
        )
        .arg(
            clap::Arg::with_name("http.host")
                .long("host")
                .value_name("HOSTNAME")
                .takes_value(true)
                .default_value("127.0.0.1")
                .help("Hostname/IP address to bind to"),
        )
        .arg(
            clap::Arg::with_name("http.port")
                .short("p")
                .long("port")
                .value_name("PORT")
                .takes_value(true)
                .default_value("5636")
                .help("Port to bind to"),
        )
        .arg(
            clap::Arg::with_name("database.elasticsearch.url")
                .short("e")
                .long("elasticsearch")
                .takes_value(true)
                .value_name("URL")
                .default_value("http://localhost:9200")
                .help("Elastic Search URL"),
        )
        .arg(
            clap::Arg::with_name("database.elasticsearch.index")
                .short("i")
                .long("index")
                .takes_value(true)
                .default_value("logstash")
                .value_name("INDEX")
                .help("Elastic Search index prefix"),
        )
        .arg(
            Arg::with_name("database.elasticsearch.no-index-suffix")
                .long("no-index-suffix")
                .takes_value(false)
                .help("Do not add a suffix to the index name"),
        )
        .arg(
            Arg::with_name("database.elasticsearch.ecs")
                .long("ecs")
                .help("Enable Elastic ECS support"),
        )
        .arg(
            Arg::with_name("database.type")
                .long("database")
                .aliases(&["datastore"])
                .takes_value(true)
                .value_name("DATABASE")
                .default_value("elasticsearch")
                .help("Database type"),
        )
        .arg(
            Arg::with_name("sqlite-filename")
                .long("sqlite-filename")
                .takes_value(true)
                .value_name("FILENAME")
                .default_value("events.sqlite")
                .help("SQLite events filename"),
        )
        .arg(
            clap::Arg::with_name("no-check-certificate")
                .short("k")
                .long("no-check-certificate")
                .help("Disable TLS certificate validation"),
        )
        .arg(
            Arg::with_name("http.request-logging")
                .long("http-request-logging")
                .env("EVEBOX_HTTP_REQUEST_LOGGING")
                .help("Enable HTTP access logging"),
        )
        .arg(
            Arg::with_name("http.tls.enabled")
                .long("tls")
                .help("Enable TLS"),
        )
        .arg(
            Arg::with_name("http.tls.certificate")
                .long("tls-cert")
                .takes_value(true)
                .value_name("FILENAME")
                .help("TLS certificate filename"),
        )
        .arg(
            Arg::with_name("http.tls.key")
                .long("tls-key")
                .takes_value(true)
                .value_name("FILENAME")
                .help("TLS key filename"),
        )
        .arg(
            Arg::with_name("input.filename")
                .long("input")
                .takes_value(true)
                .value_name("FILENAME")
                .help("Input Eve file to read"),
        )
        .arg(
            Arg::with_name("end")
                .long("end")
                .help("Read from end (tail) of input file"),
        )
        .arg(
            Arg::with_name("input-start")
                .long("input-start")
                .hidden(true),
        );

    let agent = SubCommand::with_name("agent")
        .about("EveBox Agent")
        .arg(
            Arg::with_name("config")
                .short("c")
                .long("config")
                .takes_value(true)
                .help("Configuration file"),
        )
        .arg(
            Arg::with_name("server.url")
                .long("server")
                .takes_value(true)
                .value_name("URL")
                .help("EveBox Server URL"),
        )
        .arg(Arg::with_name("geoip.enabled").long("enable-geoip"))
        .arg(
            Arg::with_name("stdout")
                .long("stdout")
                .help("Print events to stdout"),
        )
        .arg(
            Arg::with_name("bookmark-directory")
                .long("bookmark-directory")
                .takes_value(true)
                .hidden(true),
        );

    let oneshot = SubCommand::with_name("oneshot")
        .about("Import a single eve.json and review in EveBox")
        .arg(
            Arg::with_name("limit")
                .long("limit")
                .takes_value(true)
                .help("Limit the number of events read"),
        )
        .arg(
            Arg::with_name("no-open")
                .long("no-open")
                .help("Don't open browser"),
        )
        .arg(
            Arg::with_name("no-wait")
                .long("no-wait")
                .help("Don't wait for events to load"),
        )
        .arg(
            Arg::with_name("database-filename")
                .long("database-filename")
                .takes_value(true)
                .default_value("./oneshot.sqlite")
                .help("Database filename"),
        )
        // --host, but keep th name as http.host to be campatible with the
        // EVEBOX_HTTP_HOST environment variable.
        .arg(
            clap::Arg::with_name("http.host")
                .long("host")
                .value_name("HOSTNAME")
                .takes_value(true)
                .default_value("127.0.0.1")
                .help("Hostname/IP address to bind to"),
        )
        .arg(Arg::with_name("INPUT").required(true).index(1));

    let elastic_import = SubCommand::with_name("elastic-import")
        .alias("esimport")
        .about("Import to Elastic Search")
        .arg(
            Arg::with_name("config")
                .short("c")
                .long("config")
                .takes_value(true)
                .help("Configuration file"),
        )
        .arg(
            Arg::with_name("oneshot")
                .long("oneshot")
                .help("One shot mode (exit on EOF)"),
        )
        .arg(
            Arg::with_name("end")
                .long("end")
                .help("Start at end of file"),
        )
        .arg(
            clap::Arg::with_name("elasticsearch")
                .short("e")
                .long("elasticsearch")
                .takes_value(true)
                .default_value("http://localhost:9200")
                .help("Elastic Search URL"),
        )
        .arg(
            Arg::with_name("index")
                .long("index")
                .takes_value(true)
                .default_value("logstash")
                .help("Elastic Search index prefix"),
        )
        .arg(
            Arg::with_name("no-index-suffix")
                .long("no-index-suffix")
                .takes_value(false)
                .help("Do not add a suffix to the index name"),
        )
        .arg(
            Arg::with_name("bookmark")
                .long("bookmark")
                .help("Enable bookmarking"),
        )
        .arg(
            Arg::with_name("bookmark-filename")
                .long("bookmark-filename")
                .takes_value(true)
                //.default_value(".bookmark")
                .default_value("")
                .help("Bookmark filename"),
        )
        .arg(
            Arg::with_name("bookmark-dir")
                .long("bookmark-dir")
                .takes_value(true)
                .default_value(".")
                .help("Bookmark directory"),
        )
        .arg(
            Arg::with_name("stdout")
                .long("stdout")
                .help("Print events to stdout"),
        )
        .arg(
            Arg::with_name("username")
                .long("username")
                .short("u")
                .takes_value(true)
                .help("Elasticsearch username"),
        )
        .arg(
            Arg::with_name("password")
                .long("password")
                .short("p")
                .takes_value(true)
                .help("Elasticsearch password"),
        )
        .arg(
            clap::Arg::with_name(evebox::commands::elastic_import::NO_CHECK_CERTIFICATE)
                .short("k")
                .long("no-check-certificate")
                .help("Disable TLS certificate validation"),
        )
        .arg(
            Arg::with_name("geoip.disabled")
                .long("no-geoip")
                .help("Disable GeoIP"),
        )
        .arg(
            Arg::with_name("geoip.database-filename")
                .long("geoip-database")
                .takes_value(true)
                .value_name("filename")
                .help("GeoIP database filename"),
        )
        .arg(
            Arg::with_name("input")
                .multiple(true)
                .value_name("INPUT")
                .index(1),
        );

    let mut parser = parser
        .subcommand(server)
        .subcommand(elastic_import)
        .subcommand(oneshot)
        .subcommand(agent)
        .subcommand(sqlite_import)
        .subcommand(evebox::commands::config::config_subcommand())
        .subcommand(elastic_debug());

    let matches = parser.clone().get_matches();

    // Initialize logging.
    let log_level = {
        let verbosity = matches.occurrences_of("verbose");
        if verbosity > 1 {
            tracing::Level::TRACE
        } else if verbosity > 0 {
            tracing::Level::DEBUG
        } else {
            tracing::Level::INFO
        }
    };
    logger::init_logger(log_level);
    logger::init_stdlog();

    let rc: anyhow::Result<()> = match matches.subcommand() {
        ("server", Some(args)) => {
            evebox::server::main(args).await?;
            Ok(())
        }
        ("version", _) => {
            version::print_version();
            Ok(())
        }
        ("sqlite-import", Some(args)) => evebox::commands::sqlite_import::main(args).await,
        ("elastic-import", Some(args)) => {
            if let Err(err) = evebox::commands::elastic_import::main(&args).await {
                log::error!("{}", err);
                std::process::exit(1);
            }
            Ok(())
        }
        ("elastic-debug", Some(args)) => evebox::commands::elastic_debug::main(args).await,
        ("oneshot", Some(args)) => evebox::commands::oneshot::main(&args).await,
        ("agent", Some(args)) => evebox::agent::main(&args).await,
        ("config", Some(args)) => evebox::commands::config::main(&args),
        ("", None) => {
            parser.print_help().ok();
            // print_help doesn't output a new line at the end, so fix that up...
            println!();
            std::process::exit(1);
        }
        (command, _) => {
            log::error!("command \"{}\" not implemented yet", command);
            std::process::exit(1);
        }
    };
    if let Err(err) = rc {
        log::error!("error: {}", err);
        std::process::exit(1);
    } else {
        Ok(())
    }
}

fn elastic_debug() -> clap::App<'static, 'static> {
    SubCommand::with_name("elastic-debug").arg(
        Arg::with_name("elasticsearch")
            .long("elasticsearch")
            .short("e")
            .takes_value(true)
            .default_value("127.0.0.1"),
    )
}
<|endoftext|>


<commit_before>// Copyright (c) <2015> <lummax>
// Licensed under MIT (http://opensource.org/licenses/MIT)

use std::collections::VecMap;
use std::num::Int;
use std::{os, mem};

use gc_object::GCObjectRef;
use constants::{BLOCK_SIZE, LINE_SIZE, NUM_LINES_PER_BLOCK};

pub struct BlockInfo {
    mmap: os::MemoryMap,
    line_counter: VecMap<u8>,
    hole_count: u8,
    evacuation_candidate: bool,
}

impl BlockInfo {
    pub fn new(mmap: os::MemoryMap) -> BlockInfo {
        debug_assert!(mmap.len() > mem::size_of::<BlockInfo>());
        let mut block = BlockInfo {
            mmap: mmap,
            line_counter: VecMap::with_capacity(NUM_LINES_PER_BLOCK),
            hole_count: 0,
            evacuation_candidate: false,
        };
        block.clear_line_counts();
        return block;
    }

    pub fn set_evacuation_candidate(&mut self, hole_count: u8) {
        debug!("Set block {:p} to evacuation_candidate={} ({} holes)",
               &self, self.hole_count >= hole_count, self.hole_count);
        self.evacuation_candidate = self.hole_count >= hole_count;
    }

    pub fn is_evacuation_candidate(&self) -> bool{
        return self.evacuation_candidate;
    }

    pub fn increment_lines(&mut self, object: GCObjectRef) {
        self.update_line_nums(object, true);
    }

    pub fn decrement_lines(&mut self, object: GCObjectRef) {
        self.update_line_nums(object, false);
    }

    pub fn count_holes_and_marked_lines(&self) -> (u8, u8) {
        return (self.hole_count,
                self.line_counter.values().filter(|&e| *e != 0).count() as u8);
    }

    pub fn count_holes_and_available_lines(&self) -> (u8, u8) {
        return (self.hole_count,
                self.line_counter.values().filter(|&e| *e == 0).count() as u8);
    }

    pub fn clear_line_counts(&mut self) {
        for index in range(0, NUM_LINES_PER_BLOCK) {
            self.line_counter.insert(index, 0);
        }
    }

    pub fn reset(&mut self) {
        self.clear_line_counts();
        self.hole_count = 0;
        self.evacuation_candidate = false;
    }

    pub fn is_empty(&self) -> bool {
        return self.line_counter.values().all(|v| *v == 0);
    }

    pub fn offset(&self, offset: uint) -> GCObjectRef {
        return unsafe{ self.mmap.data().offset(offset as int) } as GCObjectRef;
    }

    pub fn scan_block(&self, last_high_offset: u16) -> Option<(u16, u16)> {
        debug!("Scanning block {:p} for a hole", self);
        let last_high_index = last_high_offset as uint / LINE_SIZE;
        let mut low_index = NUM_LINES_PER_BLOCK - 1;
        for index in range(last_high_index + 1, NUM_LINES_PER_BLOCK) {
            if self.line_counter.get(&index).map_or(true, |c| *c == 0) {
                // +1 to skip the next line in case an object straddles lines
                low_index = index + 1;
                break;
            }
        }
        debug!("Found low index {} in block {:p}", low_index, self);
        let mut high_index = NUM_LINES_PER_BLOCK;
        for index in range(low_index + 1, NUM_LINES_PER_BLOCK) {
            if self.line_counter.get(&index).map_or(false, |c| *c != 0) {
                high_index = index;
                break;
            }
        }
        debug!("Found high index {} in block {:p}", high_index, self);
        return if low_index < (NUM_LINES_PER_BLOCK - 1) {
            Some(((low_index * LINE_SIZE) as u16,
            (high_index * LINE_SIZE - 1) as u16))
        } else {
            debug!("Found no hole in block {:p}", self);
            None
        };
    }

    pub fn count_holes(&mut self) {
        let holes = self.line_counter.values()
            .fold((0, false), |(holes, in_hole), &elem|
                  match (in_hole, elem) {
                    (false, 0) => (holes + 1, true),
                    (_, _) => (holes, false),
                  }).0;
        self.hole_count = holes;
    }
}

impl BlockInfo {
    fn object_to_line_num(object: GCObjectRef) -> uint {
        return (object as uint % BLOCK_SIZE) / LINE_SIZE;
    }

    fn update_line_nums(&mut self, object: GCObjectRef, increment: bool) {
        // This calculates how many lines are affected starting from a
        // LINE_SIZE aligned address. So it might not mark enough lines. But
        // that does not matter as we always skip a line in scan_block()
        let line_num = BlockInfo::object_to_line_num(object);
        let object_size = unsafe{ (*object).object_size() };
        for line in range(line_num, line_num + (object_size / LINE_SIZE) + 1) {
            match increment {
                true => {
                    if self.line_counter.contains_key(&line) {
                        if let Some(val) = self.line_counter.get_mut(&line) {
                            *val += 1;
                        }
                    } else { self.line_counter.insert(line, 1); }
                    debug!("Incremented line count for line {} to {}", line,
                           self.line_counter.get(&line).unwrap());
                },
                false => {
                    if self.line_counter.contains_key(&line) {
                        if let Some(val) = self.line_counter.get_mut(&line) {
                            *val = Int::saturating_sub(*val, 1);
                        }
                    } else { self.line_counter.insert(line, 0); }
                    debug!("Decremented line count for line {} to {}", line,
                           self.line_counter.get(&line).unwrap());
                }
            }
        }
    }
}

<commit_msg>Fix BlockInfo.scanBlock()<commit_after>// Copyright (c) <2015> <lummax>
// Licensed under MIT (http://opensource.org/licenses/MIT)

use std::collections::VecMap;
use std::num::Int;
use std::{os, mem};

use gc_object::GCObjectRef;
use constants::{BLOCK_SIZE, LINE_SIZE, NUM_LINES_PER_BLOCK};

pub struct BlockInfo {
    mmap: os::MemoryMap,
    line_counter: VecMap<u8>,
    hole_count: u8,
    evacuation_candidate: bool,
}

impl BlockInfo {
    pub fn new(mmap: os::MemoryMap) -> BlockInfo {
        debug_assert!(mmap.len() > mem::size_of::<BlockInfo>());
        let mut block = BlockInfo {
            mmap: mmap,
            line_counter: VecMap::with_capacity(NUM_LINES_PER_BLOCK),
            hole_count: 0,
            evacuation_candidate: false,
        };
        block.clear_line_counts();
        return block;
    }

    pub fn set_evacuation_candidate(&mut self, hole_count: u8) {
        debug!("Set block {:p} to evacuation_candidate={} ({} holes)",
               &self, self.hole_count >= hole_count, self.hole_count);
        self.evacuation_candidate = self.hole_count >= hole_count;
    }

    pub fn is_evacuation_candidate(&self) -> bool{
        return self.evacuation_candidate;
    }

    pub fn increment_lines(&mut self, object: GCObjectRef) {
        self.update_line_nums(object, true);
    }

    pub fn decrement_lines(&mut self, object: GCObjectRef) {
        self.update_line_nums(object, false);
    }

    pub fn count_holes_and_marked_lines(&self) -> (u8, u8) {
        return (self.hole_count,
                self.line_counter.values().filter(|&e| *e != 0).count() as u8);
    }

    pub fn count_holes_and_available_lines(&self) -> (u8, u8) {
        return (self.hole_count,
                self.line_counter.values().filter(|&e| *e == 0).count() as u8);
    }

    pub fn clear_line_counts(&mut self) {
        for index in range(0, NUM_LINES_PER_BLOCK) {
            self.line_counter.insert(index, 0);
        }
    }

    pub fn reset(&mut self) {
        self.clear_line_counts();
        self.hole_count = 0;
        self.evacuation_candidate = false;
    }

    pub fn is_empty(&self) -> bool {
        return self.line_counter.values().all(|v| *v == 0);
    }

    pub fn offset(&self, offset: uint) -> GCObjectRef {
        return unsafe{ self.mmap.data().offset(offset as int) } as GCObjectRef;
    }

    pub fn scan_block(&self, last_high_offset: u16) -> Option<(u16, u16)> {
        let last_high_index = last_high_offset as uint / LINE_SIZE;
        debug!("Scanning block {:p} for a hole with last_high_offset {}",
               self, last_high_index);
        let mut low_index = NUM_LINES_PER_BLOCK - 1;
        for index in range(last_high_index + 1, NUM_LINES_PER_BLOCK) {
            if self.line_counter.get(&index).map_or(true, |c| *c == 0) {
                // +1 to skip the next line in case an object straddles lines
                low_index = index + 1;
                break;
            }
        }
        let mut high_index = NUM_LINES_PER_BLOCK;
        for index in range(low_index, NUM_LINES_PER_BLOCK) {
            if self.line_counter.get(&index).map_or(false, |c| *c != 0) {
                high_index = index;
                break;
            }
        }
        if low_index == high_index && high_index != (NUM_LINES_PER_BLOCK - 1) {
            debug!("Rescan: Found single line hole? in block {:p}", self);
            return self.scan_block((high_index * LINE_SIZE - 1) as u16);
        } else if low_index < (NUM_LINES_PER_BLOCK - 1) {
            debug!("Found low index {} and high index {} in block {:p}",
                   low_index, high_index, self);
            return Some(((low_index * LINE_SIZE) as u16,
                         (high_index * LINE_SIZE - 1) as u16));
        }
        debug!("Found no hole in block {:p}", self);
        return None;
    }

    pub fn count_holes(&mut self) {
        let holes = self.line_counter.values()
            .fold((0, false), |(holes, in_hole), &elem|
                  match (in_hole, elem) {
                    (false, 0) => (holes + 1, true),
                    (_, _) => (holes, false),
                  }).0;
        self.hole_count = holes;
    }
}

impl BlockInfo {
    fn object_to_line_num(object: GCObjectRef) -> uint {
        return (object as uint % BLOCK_SIZE) / LINE_SIZE;
    }

    fn update_line_nums(&mut self, object: GCObjectRef, increment: bool) {
        // This calculates how many lines are affected starting from a
        // LINE_SIZE aligned address. So it might not mark enough lines. But
        // that does not matter as we always skip a line in scan_block()
        let line_num = BlockInfo::object_to_line_num(object);
        let object_size = unsafe{ (*object).object_size() };
        for line in range(line_num, line_num + (object_size / LINE_SIZE) + 1) {
            match increment {
                true => {
                    if self.line_counter.contains_key(&line) {
                        if let Some(val) = self.line_counter.get_mut(&line) {
                            *val += 1;
                        }
                    } else { self.line_counter.insert(line, 1); }
                    debug!("Incremented line count for line {} to {}", line,
                           self.line_counter.get(&line).unwrap());
                },
                false => {
                    if self.line_counter.contains_key(&line) {
                        if let Some(val) = self.line_counter.get_mut(&line) {
                            *val = Int::saturating_sub(*val, 1);
                        }
                    } else { self.line_counter.insert(line, 0); }
                    debug!("Decremented line count for line {} to {}", line,
                           self.line_counter.get(&line).unwrap());
                }
            }
        }
    }
}

<|endoftext|>


<commit_before>// Implements http://rosettacode.org/wiki/Check_that_file_exists
// not_tested

fn main() {
    let paths = ["input.txt", "docs"];
    for path in paths.iter().map(|&x| Path::new(x)) {
        if path.exists() {
            println!("{} exists.", path.display());
        } else {
            println!("{} does not exist.", path.display());
        }
    }
}
<commit_msg>check file: style<commit_after>// Implements http://rosettacode.org/wiki/Check_that_file_exists
// not_tested

fn main() {
    let paths = ["input.txt", "docs"];
    for path in paths.iter().map(|&x| Path::new(x)) {
        let msg = match path.exists() {
            true => "exists",
            false => "does not exist"
        };

        println!("{} {}.", path.display(), msg);
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use serde_json;

pub mod legacy;

pub use self::legacy::*;

#[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ClaimsV2 {
    #[serde(flatten)]
    pub registered: RegisteredClaims,
    #[serde(flatten)]
    pub private: BTreeMap<String, serde_json::Value>,
}

pub type SecondsSinceEpoch = u64;

#[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct RegisteredClaims {
    #[serde(rename="iss", skip_serializing_if = "Option::is_none")]
    pub issuer: Option<String>,

    #[serde(rename="sub", skip_serializing_if = "Option::is_none")]
    pub subject: Option<String>,

    #[serde(rename="aud", skip_serializing_if = "Option::is_none")]
    pub audience: Option<String>,

    #[serde(rename="exp", skip_serializing_if = "Option::is_none")]
    pub expiration: Option<SecondsSinceEpoch>,

    #[serde(rename="nbf", skip_serializing_if = "Option::is_none")]
    pub not_before: Option<SecondsSinceEpoch>,

    #[serde(rename="nbf", skip_serializing_if = "Option::is_none")]
    pub issued_at: Option<SecondsSinceEpoch>,

    #[serde(rename="jti", skip_serializing_if = "Option::is_none")]
    pub json_web_token_id: Option<String>,
}
<commit_msg>Migrate tests<commit_after>use std::collections::BTreeMap;
use serde_json;

pub mod legacy;

pub use self::legacy::*;

#[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ClaimsV2 {
    #[serde(flatten)]
    pub registered: RegisteredClaims,
    #[serde(flatten)]
    pub private: BTreeMap<String, serde_json::Value>,
}

pub type SecondsSinceEpoch = u64;

#[derive(Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct RegisteredClaims {
    #[serde(rename="iss", skip_serializing_if = "Option::is_none")]
    pub issuer: Option<String>,

    #[serde(rename="sub", skip_serializing_if = "Option::is_none")]
    pub subject: Option<String>,

    #[serde(rename="aud", skip_serializing_if = "Option::is_none")]
    pub audience: Option<String>,

    #[serde(rename="exp", skip_serializing_if = "Option::is_none")]
    pub expiration: Option<SecondsSinceEpoch>,

    #[serde(rename="nbf", skip_serializing_if = "Option::is_none")]
    pub not_before: Option<SecondsSinceEpoch>,

    #[serde(rename="nbf", skip_serializing_if = "Option::is_none")]
    pub issued_at: Option<SecondsSinceEpoch>,

    #[serde(rename="jti", skip_serializing_if = "Option::is_none")]
    pub json_web_token_id: Option<String>,
}

#[cfg(test)]
mod tests {
    use std::default::Default;
    use crate::claims::ClaimsV2;
    use crate::Component;
    use serde_json::Value;

    // {"iss":"mikkyang.com","exp":1302319100,"custom_claim":true}
    const ENCODED_PAYLOAD: &'static str = "eyJpc3MiOiJtaWtreWFuZy5jb20iLCJleHAiOjEzMDIzMTkxMDAsImN1c3RvbV9jbGFpbSI6dHJ1ZX0K";

    #[test]
    fn registered_claims() {
        let claims = ClaimsV2::from_base64(ENCODED_PAYLOAD).unwrap();

        assert_eq!(claims.registered.issuer.unwrap(), "mikkyang.com");
        assert_eq!(claims.registered.expiration.unwrap(), 1302319100);
    }

    #[test]
    fn private_claims() {
        let claims = ClaimsV2::from_base64(ENCODED_PAYLOAD).unwrap();

        assert_eq!(claims.private["custom_claim"], Value::Bool(true));
    }

    #[test]
    fn roundtrip() {
        let mut claims: ClaimsV2 = Default::default();
        claims.registered.issuer = Some("mikkyang.com".into());
        claims.registered.expiration = Some(1302319100);
        let enc = claims.to_base64().unwrap();
        assert_eq!(claims, ClaimsV2::from_base64(&*enc).unwrap());
    }
}
<|endoftext|>


<commit_before>#![allow(unused_imports, unused_variables, unused_attributes, dead_code)]

use clap::error::{ContextKind, ContextValue, ErrorKind};
use clap::Error;
use clap::{arg, Arg, ArgMatches, Command};

use std::collections::HashMap;
use std::env;
use std::fmt::Debug;
use std::ops::Range;
use std::path::PathBuf;
use std::str::FromStr;

use crate::cfg::loader::Loader;
use crate::cfg::spec::{Nargs, Otto, Param, Spec, Task, Value};

#[macro_use]
use super::macros;
use super::error;

static OTTOFILE: &str = "./otto.yml";

fn print_type_of<T: ?Sized>(t: &T)
where
    T: Debug,
{
    println!("type={} value={:#?}", std::any::type_name::<T>(), t);
}

fn extract(item: (ContextKind, &ContextValue)) -> Option<&ContextValue> {
    let (k, v) = item;
    if k == ContextKind::InvalidArg {
        return Some(v);
    }
    None
}

pub trait GetKnownMatches {
    fn get_known_matches(&self) -> Result<(ArgMatches, Vec<String>), Error>;
    fn get_known_matches_from(&self, args: &mut Vec<String>) -> Result<(ArgMatches, Vec<String>), Error>;
}

impl<'a> GetKnownMatches for Command<'a> {
    fn get_known_matches(&self) -> Result<(ArgMatches, Vec<String>), Error> {
        let mut args: Vec<String> = env::args().collect();
        self.get_known_matches_from(&mut args)
    }
    fn get_known_matches_from(&self, args: &mut Vec<String>) -> Result<(ArgMatches, Vec<String>), Error> {
        let mut rem: Vec<String> = vec![];
        loop {
            match self.clone().try_get_matches_from(&*args) {
                Ok(matches) => {
                    return Ok((matches, rem));
                }
                Err(error) => match error.kind() {
                    ErrorKind::UnknownArgument => {
                        let items = error.context().find_map(extract);
                        match items {
                            Some(ContextValue::String(s)) => {
                                rem.push(s.to_owned());
                                args.retain(|a| a != s);
                            }
                            Some(&_) => {
                                return Err(error);
                            }
                            None => {
                                return Err(error);
                            }
                        }
                    }
                    _ => {
                        return Err(error);
                    }
                },
            }
        }
    }
}
#[derive(Debug, Default, PartialEq, Clone)]
struct PartitionedArgs {
    args: Vec<String>,
    partitions: Vec<Range<usize>>,
}
impl PartitionedArgs {
    fn new(tasknames: &[&str]) -> Self {
        let args: Vec<String> = env::args().collect();
        let mut beg = 0;
        let mut partitions = vec![];
        for (i, arg) in args.iter().skip(1).enumerate() {
            if tasknames.iter().any(|t| t == arg) {
                let end = i + 1;
                partitions.push(Range::<usize> { start: beg, end });
                beg = end;
            }
        }
        partitions.push(Range::<usize> {
            start: beg,
            end: args.len(),
        });
        Self { args, partitions }
    }
    fn partitions(&self) -> Vec<&[String]> {
        self.partitions.iter().map(|p| &self.args[p.clone()]).collect()
    }
    fn partition(&self, index: usize) -> Option<&[String]> {
        if index < self.len() {
            return Some(&self.args[self.partitions[index].clone()]);
        }
        None
    }
    fn len(&self) -> usize {
        self.partitions.len()
    }
}

#[derive(Debug, PartialEq)]
pub struct Parser {
    ottofile: String,
}

impl Default for Parser {
    fn default() -> Self {
        Self::new()
    }
}

impl Parser {
    pub fn new() -> Self {
        Self {
            ottofile: env::var("OTTOFILE").unwrap_or_else(|_| OTTOFILE.to_owned()),
        }
    }
    fn otto_seed(&self, nerfed: bool) -> Command {
        Command::new("otto")
            .disable_help_flag(nerfed)
            .disable_version_flag(nerfed)
            .arg(
                Arg::new("ottofile")
                    .takes_value(true)
                    .short('o')
                    .long("ottofile")
                    .help("override default path to ottofile"),
            )
    }

    pub fn divine_ottofile(&self) -> PathBuf {
        let ottofile = match self.otto_seed(true).get_known_matches() {
            Ok((matches, _)) => match matches.value_of("ottofile").map(str::to_string) {
                Some(s) => s,
                None => self.ottofile.clone(),
            },
            Err(error) => self.ottofile.clone(),
        };
        ottofile.into()
    }
    pub fn parse(&self) -> Vec<ArgMatches> {
        let ottofile = self.divine_ottofile();
        if ottofile.exists() {
            let loader = Loader::new(&ottofile);
            let spec = loader.load().unwrap();
            let task_names = spec.otto.task_names();
            let mut commands = HashMap::<String, Command>::new();
            for task_name in task_names.clone() {
                let task = &spec.otto.tasks[<&str>::clone(&task_name)];
                let command = Parser::task_to_command(task);
                *commands.get_mut(<&str>::clone(&task_name)).unwrap() = command;
            }
            let pa = PartitionedArgs::new(&task_names.clone());
        } else {
            let after_help = format!("ottofile={:?} does not exist!", ottofile);
            let otto = self
                .otto_seed(false)
                .arg_required_else_help(true)
                .after_help(after_help.as_str());
            let matches = otto.get_matches_from(vec!["--help"]);
        }
        vec![]
    }
    fn task_to_command(task: &Task) -> Command {
        let mut command = Command::new(&task.name);
        for param in task.params.values() {
            command = command.arg(Parser::param_to_arg(param));
        }
        command
    }
    fn param_to_arg(param: &Param) -> Arg {
        let mut arg = Arg::new(&*param.name);
        if let Some(short) = &param.short {
            arg = arg.short(short.chars().next().unwrap());
        }
        if let Some(long) = &param.long {
            arg = arg.long(long);
        }
        arg
    }
}
<commit_msg>fixedup and used &str instead<commit_after>#![allow(unused_imports, unused_variables, unused_attributes, dead_code)]

use clap::error::{ContextKind, ContextValue, ErrorKind};
use clap::Error;
use clap::{arg, Arg, ArgMatches, Command};

use std::collections::HashMap;
use std::env;
use std::fmt::Debug;
use std::ops::Range;
use std::path::PathBuf;
use std::str::FromStr;

use crate::cfg::loader::Loader;
use crate::cfg::spec::{Nargs, Otto, Param, Spec, Task, Value};

#[macro_use]
use super::macros;
use super::error;

static OTTOFILE: &str = "./otto.yml";

fn print_type_of<T: ?Sized>(t: &T)
where
    T: Debug,
{
    println!("type={} value={:#?}", std::any::type_name::<T>(), t);
}

fn extract(item: (ContextKind, &ContextValue)) -> Option<&ContextValue> {
    let (k, v) = item;
    if k == ContextKind::InvalidArg {
        return Some(v);
    }
    None
}

pub trait GetKnownMatches {
    fn get_known_matches(&self) -> Result<(ArgMatches, Vec<String>), Error>;
    fn get_known_matches_from(&self, args: &mut Vec<String>) -> Result<(ArgMatches, Vec<String>), Error>;
}

impl<'a> GetKnownMatches for Command<'a> {
    fn get_known_matches(&self) -> Result<(ArgMatches, Vec<String>), Error> {
        let mut args: Vec<String> = env::args().collect();
        self.get_known_matches_from(&mut args)
    }
    fn get_known_matches_from(&self, args: &mut Vec<String>) -> Result<(ArgMatches, Vec<String>), Error> {
        let mut rem: Vec<String> = vec![];
        loop {
            match self.clone().try_get_matches_from(&*args) {
                Ok(matches) => {
                    return Ok((matches, rem));
                }
                Err(error) => match error.kind() {
                    ErrorKind::UnknownArgument => {
                        let items = error.context().find_map(extract);
                        match items {
                            Some(ContextValue::String(s)) => {
                                rem.push(s.to_owned());
                                args.retain(|a| a != s);
                            }
                            Some(&_) => {
                                return Err(error);
                            }
                            None => {
                                return Err(error);
                            }
                        }
                    }
                    _ => {
                        return Err(error);
                    }
                },
            }
        }
    }
}
#[derive(Debug, Default, PartialEq, Clone)]
struct PartitionedArgs {
    args: Vec<String>,
    partitions: Vec<Range<usize>>,
}
impl PartitionedArgs {
    fn new(tasknames: &[&str]) -> Self {
        let args: Vec<String> = env::args().collect();
        let mut beg = 0;
        let mut partitions = vec![];
        for (i, arg) in args.iter().skip(1).enumerate() {
            if tasknames.iter().any(|t| t == arg) {
                let end = i + 1;
                partitions.push(Range::<usize> { start: beg, end });
                beg = end;
            }
        }
        partitions.push(Range::<usize> {
            start: beg,
            end: args.len(),
        });
        Self { args, partitions }
    }
    fn partitions(&self) -> Vec<&[String]> {
        self.partitions.iter().map(|p| &self.args[p.clone()]).collect()
    }
    fn partition(&self, index: usize) -> Option<&[String]> {
        if index < self.len() {
            return Some(&self.args[self.partitions[index].clone()]);
        }
        None
    }
    fn len(&self) -> usize {
        self.partitions.len()
    }
}

#[derive(Debug, PartialEq)]
pub struct Parser {
    ottofile: String,
}

impl Default for Parser {
    fn default() -> Self {
        Self::new()
    }
}

impl Parser {
    pub fn new() -> Self {
        Self {
            ottofile: env::var("OTTOFILE").unwrap_or_else(|_| OTTOFILE.to_owned()),
        }
    }
    fn otto_seed(&self, nerfed: bool) -> Command {
        Command::new("otto")
            .disable_help_flag(nerfed)
            .disable_version_flag(nerfed)
            .arg(
                Arg::new("ottofile")
                    .takes_value(true)
                    .short('o')
                    .long("ottofile")
                    .help("override default path to ottofile"),
            )
    }

    pub fn divine_ottofile(&self) -> PathBuf {
        let ottofile = match self.otto_seed(true).get_known_matches() {
            Ok((matches, _)) => match matches.value_of("ottofile").map(str::to_string) {
                Some(s) => s,
                None => self.ottofile.clone(),
            },
            Err(error) => self.ottofile.clone(),
        };
        ottofile.into()
    }
    pub fn parse(&self) -> Vec<ArgMatches> {
        let ottofile = self.divine_ottofile();
        if ottofile.exists() {
            let loader = Loader::new(&ottofile);
            let spec = loader.load().unwrap();
            let task_names = spec.otto.task_names();
            let mut commands = HashMap::<&str, Command>::new();
            for task_name in task_names.clone() {
                let task = &spec.otto.tasks[task_name];
                let command = Parser::task_to_command(task);
                commands.insert(task_name, command);
            }
            let pa = PartitionedArgs::new(&task_names.clone());
        } else {
            let after_help = format!("ottofile={:?} does not exist!", ottofile);
            let otto = self
                .otto_seed(false)
                .arg_required_else_help(true)
                .after_help(after_help.as_str());
            let matches = otto.get_matches_from(vec!["--help"]);
        }
        vec![]
    }
    fn task_to_command(task: &Task) -> Command {
        let mut command = Command::new(&task.name);
        for param in task.params.values() {
            command = command.arg(Parser::param_to_arg(param));
        }
        command
    }
    fn param_to_arg(param: &Param) -> Arg {
        let mut arg = Arg::new(&*param.name);
        if let Some(short) = &param.short {
            arg = arg.short(short.chars().next().unwrap());
        }
        if let Some(long) = &param.long {
            arg = arg.long(long);
        }
        arg
    }
}
<|endoftext|>


<commit_before>//! Signals are a list of methods associated with a name that will be called
//! when the signal is triggered.
//!
//! Signals are stored with the object in its metatable,
//! the methods defined here are just to make it easier to use.

use rlua::{self, Function, Lua, Table, ToLua, ToLuaMulti, Value};

use ::GLOBAL_SIGNALS;
use common::object::{Object, ObjectStateType};

/// Connects functions to a signal. Creates a new entry in the table if it
/// doesn't exist.
pub fn connect_signal<S: ObjectStateType>(lua: &Lua,
                                          obj: Object<S>,
                                          name: String,
                                          funcs: &[Function])
                                          -> rlua::Result<()> {
    let signals = obj.signals()?;
    connect_signals(lua, signals, name, funcs)
}

pub fn connect_signals<'lua>(lua: &'lua Lua,
                             signals: Table<'lua>,
                             name: String,
                             funcs: &[Function])
                             -> rlua::Result<()> {
    if let Ok(Value::Table(table)) = signals.get::<_, Value>(name.as_str()) {
        let mut length = table.len()? + 1;
        for func in funcs {
            table.set(length, func.clone())?;
            length += 1;
        }
        Ok(())
    } else {
        let table = lua.create_table()?;
        for (index, func) in funcs.into_iter().enumerate() {
            table.set(index + 1, func.clone())?;
        }
        signals.set(name, table)
    }
}

pub fn disconnect_signal<S: ObjectStateType>(lua: &Lua,
                                             obj: Object<S>,
                                             name: String,
)
                                             -> rlua::Result<()> {
    let signals = obj.signals()?;
    disconnect_signals(lua, signals, name)
}

pub fn disconnect_signals(_: &Lua, signals: Table, name: String) -> rlua::Result<()> {
    signals.set(name, Value::Nil)
}

/// Evaluate the functions associated with a signal.
pub fn emit_object_signal<'lua, A, S>(lua: &'lua Lua,
                                      obj: Object<'lua, S>,
                                      name: String,
                                      args: A)
                                      -> rlua::Result<()>
    where A: ToLuaMulti<'lua> + Clone,
          S: ObjectStateType
{
    let signals = obj.signals()?;
    let mut args = args.to_lua_multi(lua)?;
    args.push_front(obj.to_lua(lua)?);
    emit_signals(lua, signals, name, args)
}

fn emit_signals<'lua, A>(_: &'lua Lua,
                         signals: Table<'lua>,
                         name: String,
                         args: A)
                         -> rlua::Result<()>
    where A: ToLuaMulti<'lua> + Clone
{
    if let Ok(Value::Table(table)) = signals.get::<_, Value>(name.clone()) {
        for entry in table.pairs::<Value, Function>() {
            if let Ok((_, func)) = entry {
                match func.call(args.clone()) {
                    Ok(()) => {}
                    Err(e) => {
                        error!("Error while emitting signal {}: {}", name, e);
                    }
                };
            }
        }
    }
    Ok(())
}

/// Connect the function to the named signal in the global signal list.
pub fn global_connect_signal<'lua>(lua: &'lua Lua,
                                   (name, func): (String, Function<'lua>))
                                   -> rlua::Result<()> {
    let global_signals = lua.named_registry_value::<Table>(GLOBAL_SIGNALS)?;
    connect_signals(lua, global_signals, name, &[func])
}

/// Disconnect the function from the named signal in the global signal list.
pub fn global_disconnect_signal<'lua>(lua: &'lua Lua, name: String) -> rlua::Result<()> {
    let global_signals = lua.named_registry_value::<Table>(GLOBAL_SIGNALS)?;
    disconnect_signals(lua, global_signals, name)
}

/// Emit the signal with the given name from the global signal list.
pub fn global_emit_signal<'lua>(lua: &'lua Lua, (name, args): (String, Value)) -> rlua::Result<()> {
    let global_signals = lua.named_registry_value::<Table>(GLOBAL_SIGNALS)?;
    emit_signals(lua, global_signals, name, args)
}
<commit_msg>Expose emit_signals<commit_after>//! Signals are a list of methods associated with a name that will be called
//! when the signal is triggered.
//!
//! Signals are stored with the object in its metatable,
//! the methods defined here are just to make it easier to use.

use rlua::{self, Function, Lua, Table, ToLua, ToLuaMulti, Value};

use ::GLOBAL_SIGNALS;
use common::object::{Object, ObjectStateType};

/// Connects functions to a signal. Creates a new entry in the table if it
/// doesn't exist.
pub fn connect_signal<S: ObjectStateType>(lua: &Lua,
                                          obj: Object<S>,
                                          name: String,
                                          funcs: &[Function])
                                          -> rlua::Result<()> {
    let signals = obj.signals()?;
    connect_signals(lua, signals, name, funcs)
}

pub fn connect_signals<'lua>(lua: &'lua Lua,
                             signals: Table<'lua>,
                             name: String,
                             funcs: &[Function])
                             -> rlua::Result<()> {
    if let Ok(Value::Table(table)) = signals.get::<_, Value>(name.as_str()) {
        let mut length = table.len()? + 1;
        for func in funcs {
            table.set(length, func.clone())?;
            length += 1;
        }
        Ok(())
    } else {
        let table = lua.create_table()?;
        for (index, func) in funcs.into_iter().enumerate() {
            table.set(index + 1, func.clone())?;
        }
        signals.set(name, table)
    }
}

pub fn disconnect_signal<S: ObjectStateType>(lua: &Lua,
                                             obj: Object<S>,
                                             name: String,
)
                                             -> rlua::Result<()> {
    let signals = obj.signals()?;
    disconnect_signals(lua, signals, name)
}

pub fn disconnect_signals(_: &Lua, signals: Table, name: String) -> rlua::Result<()> {
    signals.set(name, Value::Nil)
}

/// Evaluate the functions associated with a signal.
pub fn emit_object_signal<'lua, A, S>(lua: &'lua Lua,
                                      obj: Object<'lua, S>,
                                      name: String,
                                      args: A)
                                      -> rlua::Result<()>
    where A: ToLuaMulti<'lua> + Clone,
          S: ObjectStateType
{
    let signals = obj.signals()?;
    let mut args = args.to_lua_multi(lua)?;
    args.push_front(obj.to_lua(lua)?);
    emit_signals(lua, signals, name, args)
}

pub fn emit_signals<'lua, A>(_: &'lua Lua,
                         signals: Table<'lua>,
                         name: String,
                         args: A)
                         -> rlua::Result<()>
    where A: ToLuaMulti<'lua> + Clone
{
    if let Ok(Value::Table(table)) = signals.get::<_, Value>(name.clone()) {
        for entry in table.pairs::<Value, Function>() {
            if let Ok((_, func)) = entry {
                match func.call(args.clone()) {
                    Ok(()) => {}
                    Err(e) => {
                        error!("Error while emitting signal {}: {}", name, e);
                    }
                };
            }
        }
    }
    Ok(())
}

/// Connect the function to the named signal in the global signal list.
pub fn global_connect_signal<'lua>(lua: &'lua Lua,
                                   (name, func): (String, Function<'lua>))
                                   -> rlua::Result<()> {
    let global_signals = lua.named_registry_value::<Table>(GLOBAL_SIGNALS)?;
    connect_signals(lua, global_signals, name, &[func])
}

/// Disconnect the function from the named signal in the global signal list.
pub fn global_disconnect_signal<'lua>(lua: &'lua Lua, name: String) -> rlua::Result<()> {
    let global_signals = lua.named_registry_value::<Table>(GLOBAL_SIGNALS)?;
    disconnect_signals(lua, global_signals, name)
}

/// Emit the signal with the given name from the global signal list.
pub fn global_emit_signal<'lua>(lua: &'lua Lua, (name, args): (String, Value)) -> rlua::Result<()> {
    let global_signals = lua.named_registry_value::<Table>(GLOBAL_SIGNALS)?;
    emit_signals(lua, global_signals, name, args)
}
<|endoftext|>


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

//! The client API itself.
//!
//! This contains three main functions that return three specific types of client:
//!
//! * `connect` returns a pair of `Stream` and `Sink`, clients can write RESP messages to the
//! `Sink` and read RESP messages from the `Stream`. Pairing requests to responses is up to the
//! client.  This is intended to be a low-level interface from which more user-friendly interfaces
//! can be built.
//! * `paired_connect` is used for most of the standard Redis commands, where one request results
//! in one response.
//! * `pubsub_connect` is used for Redis's PUBSUB functionality.

pub mod connect;
#[macro_use]
pub mod paired;
pub mod pubsub;

pub use self::connect::{connect, ClientConnection};
pub use self::paired::{paired_connect, PairedConnection};

// TODO - uncomment
// pub use self::pubsub::{pubsub_connect, PubsubConnection};

#[cfg(test)]
mod test {
    use std::fmt;
    use std::io;

    use futures::future;
    use futures::sync::oneshot;
    use futures::{stream, Future, Sink, Stream};

    use tokio::executor::current_thread;

    use resp;

    fn extract_result<F, R, E>(f: F) -> R
    where
        R: 'static,
        F: Future<Item = R, Error = E> + 'static,
        E: fmt::Debug + 'static,
    {
        let r = current_thread::run(|_| {
            let (tx, rx) = oneshot::channel();
            current_thread::spawn(f.then(move |r| match tx.send(r) {
                Ok(_) => future::ok(()),
                Err(_) => future::err(()),
            }));
            rx
        });
        r.wait()
            .expect("Result was cancelled")
            .expect("Future failed")
    }

    #[test]
    fn can_connect() {
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connection = super::connect(&addr)
            .map_err(|e| e.into())
            .and_then(|connection| {
                let a = connection
                    .sender
                    .send(resp_array!["PING", "TEST"])
                    .map_err(|e| e.into());
                let b = connection.receiver.take(1).collect();
                a.join(b)
            });

        let (_, values) = extract_result(connection);

        assert_eq!(values.len(), 1);
        assert_eq!(values[0], "TEST".into());
    }

    #[test]
    fn complex_test() {
        let addr = "127.0.0.1:6379".parse().unwrap();
        let connection = super::connect(&addr)
            .map_err(|e| e.into())
            .and_then(|connection| {
                let mut ops = Vec::<resp::RespValue>::new();
                ops.push(resp_array!["FLUSH"]);
                ops.extend(
                    (0..1000).map(|i| resp_array!["SADD", "test_set", format!("VALUE: {}", i)]),
                );
                ops.push(resp_array!["SMEMBERS", "test_set"]);
                let send = connection
                    .sender
                    .send_all(stream::iter_ok::<_, io::Error>(ops))
                    .map_err(|e| e.into());
                let receive = connection.receiver.skip(1001).take(1).collect();
                send.join(receive)
            });
        let (_, values) = extract_result(connection);
        assert_eq!(values.len(), 1);
        let values = match &values[0] {
            &resp::RespValue::Array(ref values) => values.clone(),
            _ => panic!("Not an array"),
        };
        assert_eq!(values.len(), 1000);
    }

    #[test]
    fn can_paired_connect() {
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connect_f =
            super::paired_connect(&addr, current_thread::task_executor()).and_then(|connection| {
                let res_f = connection.send(resp_array!["PING", "TEST"]).map(|v| {
                    println!("FIRST: {:?}", v);
                    v
                });
                faf!(connection.send(resp_array!["SET", "X", "123"]));
                let wait_f = connection.send(resp_array!["GET", "X"]).map(|v| {
                    println!("THIRD: {:?}", v);
                    v
                });
                res_f.join(wait_f)
            });
        let (result_1, result_2): (String, String) = extract_result(connect_f);
        assert_eq!(result_1, "TEST");
        assert_eq!(result_2, "123");
    }

    // #[test]
    // fn complex_paired_connect() {
    //     let addr = "127.0.0.1:6379".parse().unwrap();

    //     let connect_f = super::paired_connect(&addr).and_then(|connection| {
    //         connection
    //             .send(resp_array!["INCR", "CTR"])
    //             .and_then(move |value: String| {
    //                 connection.send(resp_array!["SET", "LASTCTR", value])
    //             })
    //     });
    //     let result: String = core.run(connect_f).unwrap();
    //     assert_eq!(result, "OK");
    // }

    // #[test]
    // fn sending_a_lot_of_data_test() {
    //     let addr = "127.0.0.1:6379".parse().unwrap();

    //     let test_f = super::paired_connect(&addr);
    //     let send_data = test_f.and_then(|connection| {
    //         let mut futures = Vec::with_capacity(1000);
    //         for i in 0..1000 {
    //             let key = format!("X_{}", i);
    //             faf!(connection.send(resp_array!["SET", &key, i.to_string()]));
    //             futures.push(connection.send(resp_array!["GET", key]));
    //         }
    //         futures.remove(999)
    //     });
    //     let result: String = core.run(send_data).unwrap();
    //     assert_eq!(result, "999");
    // }

    // #[test]
    // fn pubsub_test() {
    //     let addr = "127.0.0.1:6379".parse().unwrap();
    //     let paired_c = super::paired_connect(&addr);
    //     let pubsub_c = super::pubsub_connect(&addr);
    //     let msgs = paired_c.join(pubsub_c).and_then(|(paired, pubsub)| {
    //         let subscribe = pubsub.subscribe("test-topic");
    //         subscribe.and_then(move |msgs| {
    //             faf!(paired.send(resp_array!["PUBLISH", "test-topic", "test-message"]));
    //             faf!(paired.send(resp_array!["PUBLISH", "test-not-topic", "test-message-1.5"]));
    //             paired
    //                 .send(resp_array!["PUBLISH", "test-topic", "test-message2"])
    //                 .map(|_: resp::RespValue| msgs)
    //         })
    //     });
    //     let tst = msgs.and_then(|msgs| {
    //         msgs.take(2)
    //             .collect()
    //             .map_err(|_| error::internal("unreachable"))
    //     });
    //     let result = core.run(tst).unwrap();
    //     assert_eq!(result.len(), 2);
    //     assert_eq!(result[0], "test-message".into());
    //     assert_eq!(result[1], "test-message2".into());
    // }
}
<commit_msg>More re-enabled tests<commit_after>/*
 * Copyright 2017-2018 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

//! The client API itself.
//!
//! This contains three main functions that return three specific types of client:
//!
//! * `connect` returns a pair of `Stream` and `Sink`, clients can write RESP messages to the
//! `Sink` and read RESP messages from the `Stream`. Pairing requests to responses is up to the
//! client.  This is intended to be a low-level interface from which more user-friendly interfaces
//! can be built.
//! * `paired_connect` is used for most of the standard Redis commands, where one request results
//! in one response.
//! * `pubsub_connect` is used for Redis's PUBSUB functionality.

pub mod connect;
#[macro_use]
pub mod paired;
pub mod pubsub;

pub use self::connect::{connect, ClientConnection};
pub use self::paired::{paired_connect, PairedConnection};

// TODO - uncomment
// pub use self::pubsub::{pubsub_connect, PubsubConnection};

#[cfg(test)]
mod test {
    use std::fmt;
    use std::io;

    use futures::future;
    use futures::sync::oneshot;
    use futures::{stream, Future, Sink, Stream};

    use tokio::executor::current_thread;

    use resp;

    fn extract_result<F, R, E>(f: F) -> R
    where
        R: 'static,
        F: Future<Item = R, Error = E> + 'static,
        E: fmt::Debug + 'static,
    {
        let r = current_thread::run(|_| {
            let (tx, rx) = oneshot::channel();
            current_thread::spawn(f.then(move |r| match tx.send(r) {
                Ok(_) => future::ok(()),
                Err(_) => future::err(()),
            }));
            rx
        });
        r.wait()
            .expect("Result was cancelled")
            .expect("Future failed")
    }

    #[test]
    fn can_connect() {
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connection = super::connect(&addr)
            .map_err(|e| e.into())
            .and_then(|connection| {
                let a = connection
                    .sender
                    .send(resp_array!["PING", "TEST"])
                    .map_err(|e| e.into());
                let b = connection.receiver.take(1).collect();
                a.join(b)
            });

        let (_, values) = extract_result(connection);

        assert_eq!(values.len(), 1);
        assert_eq!(values[0], "TEST".into());
    }

    #[test]
    fn complex_test() {
        let addr = "127.0.0.1:6379".parse().unwrap();
        let connection = super::connect(&addr)
            .map_err(|e| e.into())
            .and_then(|connection| {
                let mut ops = Vec::<resp::RespValue>::new();
                ops.push(resp_array!["FLUSH"]);
                ops.extend(
                    (0..1000).map(|i| resp_array!["SADD", "test_set", format!("VALUE: {}", i)]),
                );
                ops.push(resp_array!["SMEMBERS", "test_set"]);
                let send = connection
                    .sender
                    .send_all(stream::iter_ok::<_, io::Error>(ops))
                    .map_err(|e| e.into());
                let receive = connection.receiver.skip(1001).take(1).collect();
                send.join(receive)
            });
        let (_, values) = extract_result(connection);
        assert_eq!(values.len(), 1);
        let values = match &values[0] {
            &resp::RespValue::Array(ref values) => values.clone(),
            _ => panic!("Not an array"),
        };
        assert_eq!(values.len(), 1000);
    }

    #[test]
    fn can_paired_connect() {
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connect_f =
            super::paired_connect(&addr, current_thread::task_executor()).and_then(|connection| {
                let res_f = connection.send(resp_array!["PING", "TEST"]).map(|v| {
                    println!("FIRST: {:?}", v);
                    v
                });
                faf!(connection.send(resp_array!["SET", "X", "123"]));
                let wait_f = connection.send(resp_array!["GET", "X"]).map(|v| {
                    println!("THIRD: {:?}", v);
                    v
                });
                res_f.join(wait_f)
            });
        let (result_1, result_2): (String, String) = extract_result(connect_f);
        assert_eq!(result_1, "TEST");
        assert_eq!(result_2, "123");
    }

    #[test]
    fn complex_paired_connect() {
        let addr = "127.0.0.1:6379".parse().unwrap();

        let connect_f =
            super::paired_connect(&addr, current_thread::task_executor()).and_then(|connection| {
                connection
                    .send(resp_array!["INCR", "CTR"])
                    .and_then(move |value: String| {
                        connection.send(resp_array!["SET", "LASTCTR", value])
                    })
            });
        let result: String = extract_result(connect_f);
        assert_eq!(result, "OK");
    }

    #[test]
    fn sending_a_lot_of_data_test() {
        let addr = "127.0.0.1:6379".parse().unwrap();

        let test_f = super::paired_connect(&addr, current_thread::task_executor());
        let send_data = test_f.and_then(|connection| {
            let mut futures = Vec::with_capacity(1000);
            for i in 0..1000 {
                let key = format!("X_{}", i);
                faf!(connection.send(resp_array!["SET", &key, i.to_string()]));
                futures.push(connection.send(resp_array!["GET", key]));
            }
            futures.remove(999)
        });
        let result: String = extract_result(send_data);
        assert_eq!(result, "999");
    }

    // #[test]
    // fn pubsub_test() {
    //     let addr = "127.0.0.1:6379".parse().unwrap();
    //     let paired_c = super::paired_connect(&addr);
    //     let pubsub_c = super::pubsub_connect(&addr);
    //     let msgs = paired_c.join(pubsub_c).and_then(|(paired, pubsub)| {
    //         let subscribe = pubsub.subscribe("test-topic");
    //         subscribe.and_then(move |msgs| {
    //             faf!(paired.send(resp_array!["PUBLISH", "test-topic", "test-message"]));
    //             faf!(paired.send(resp_array!["PUBLISH", "test-not-topic", "test-message-1.5"]));
    //             paired
    //                 .send(resp_array!["PUBLISH", "test-topic", "test-message2"])
    //                 .map(|_: resp::RespValue| msgs)
    //         })
    //     });
    //     let tst = msgs.and_then(|msgs| {
    //         msgs.take(2)
    //             .collect()
    //             .map_err(|_| error::internal("unreachable"))
    //     });
    //     let result = core.run(tst).unwrap();
    //     assert_eq!(result.len(), 2);
    //     assert_eq!(result[0], "test-message".into());
    //     assert_eq!(result[1], "test-message2".into());
    // }
}
<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    1 as c_int

    // let mut r_address;
    // unsafe {
    //     r_address = CStr::from_ptr(address);
    // }
    // let s_address = r_address.to_bytes();
    // let host_address = match str::from_utf8(s_address) {
    //     Ok(safe_str) => safe_str,
    //     Err(_) => {
    //         println!("Invalid host address");
    //         return -1 as c_int;
    //     }
    // };
    //
    // // Create and register a way to kill this client
    // let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    // let kill_tx = k_tx.clone();
    // let mut k_tx_ptr = Box::new(k_tx);
    // unsafe {
    //     register_stop_tx(&mut *k_tx_ptr);
    // }
    //
    // // Writer thread's channel
    // let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    // let mut w_tx_ptr = Box::new(w_tx);
    // unsafe {
    //     register_writer_tx(&mut *w_tx_ptr);
    // }
    //
    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();
    //
    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();
    //
    // // Exit out in standard C fashion
    // 0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<commit_msg>Debugging stuff<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    unsafe {
        on_connect_handler();
        on_disconnect_handler();
    }

    1 as c_int

    // let mut r_address;
    // unsafe {
    //     r_address = CStr::from_ptr(address);
    // }
    // let s_address = r_address.to_bytes();
    // let host_address = match str::from_utf8(s_address) {
    //     Ok(safe_str) => safe_str,
    //     Err(_) => {
    //         println!("Invalid host address");
    //         return -1 as c_int;
    //     }
    // };
    //
    // // Create and register a way to kill this client
    // let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    // let kill_tx = k_tx.clone();
    // let mut k_tx_ptr = Box::new(k_tx);
    // unsafe {
    //     register_stop_tx(&mut *k_tx_ptr);
    // }
    //
    // // Writer thread's channel
    // let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    // let mut w_tx_ptr = Box::new(w_tx);
    // unsafe {
    //     register_writer_tx(&mut *w_tx_ptr);
    // }
    //
    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();
    //
    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();
    //
    // // Exit out in standard C fashion
    // 0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<|endoftext|>


<commit_before>#![feature(proc_macro)]
#![feature(plugin)]
#![plugin(docopt_macros)]
#[macro_use] extern crate diesel;
#[macro_use] extern crate diesel_codegen;
extern crate docopt;
extern crate dotenv;
#[macro_use] extern crate lazy_static;
extern crate rand;
extern crate regex;
extern crate rustc_serialize;

pub mod counts;
pub mod database;
pub mod schema;
pub mod models;
pub mod quotes;

use self::database::{establish_connection};
use self::quotes::{retrieve_quote, store_quotes, RetrievalRequest};
use self::counts::{count_authors, count_quotes};
use std::fs::File;
use std::io::{BufReader, BufRead};
use std::path::Path;

docopt!(Args derive Debug, "
Quote Storage & Retrieval Utilities

Usage:
  quote_storage count [authors | unretrieved-quotes | retrieved-quotes]
  quote_storage store <file> [--mark-retrieved]
  quote_storage retrieve <quote-id> [--mark-retrieved]
  quote_storage retrieve [--random-retrieved | --random-unretrieved | --first-unretrieved] [--mark-retrieved]
  quote_storage (-h | --help)
  quote_storage --version

Options:
  -h --help             Show this screen.
  --version             Show version.
  --random              Retrieves a random quote.
  --random-retrieved    Retrieves a random already retrieved quote.
  --random-unretrieved  Retrieves a random quote that has not already been retrieved.
  --mark-retrieved      Marks quotes as retrieved.
");

fn main() {
    let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());

    if args.flag_version {
        let version = env!("CARGO_PKG_VERSION");

        return println!("quote_storage v{}", version);
    }

    let dbsession = establish_connection();

    // dbsession.begin_test_transaction().unwrap(); // Testing

    if args.cmd_store {
        let path = Path::new(&args.arg_file);
        let file = match File::open(&path) {
            Ok(f) => f,
            Err(e) => panic!("FIXME: Invalid file reference found :'( {}", e)
        };

        store_quotes(dbsession, &quotes_from_buffered_reader(BufReader::new(file)), args.flag_mark_retrieved);

    } else if args.cmd_retrieve {
        let request = if args.flag_random_retrieved {
            RetrievalRequest::RandomRetrieved
        } else if args.flag_random_unretrieved {
            RetrievalRequest::RandomUnretrievedAndMark(args.flag_mark_retrieved)
        } else if args.flag_first_unretrieved {
            RetrievalRequest::FirstUnretrievedAndMark(args.flag_mark_retrieved)
        } else if args.arg_quote_id.len() > 0 {
            let id = args.arg_quote_id.parse::<i64>().unwrap(); // FIXME: Handle error gracefully

            RetrievalRequest::IdAndMark(id, args.flag_mark_retrieved)
        } else {
            RetrievalRequest::Random
        };

        let quote = retrieve_quote(dbsession, request);

        println!("{}", quote);
    } else if args.cmd_count {
        let count = if args.cmd_authors {
            count_authors(dbsession)
        } else if args.cmd_unretrieved_quotes {
            count_quotes(dbsession, true)
        } else {
            count_quotes(dbsession, false)
        };

        print!("{}", count);
    } else {
        unreachable!("Should not be able to get here!")
    }
}

fn quotes_from_buffered_reader<B: BufRead>(bufreader: B) -> Vec<String> {
    bufreader.lines().map(|l| l.expect("Could not parse line")).collect()
}
<commit_msg>Cleaned up some code<commit_after>#![feature(proc_macro)]
#![feature(plugin)]
#![plugin(docopt_macros)]
#[macro_use] extern crate diesel;
#[macro_use] extern crate diesel_codegen;
extern crate docopt;
extern crate dotenv;
#[macro_use] extern crate lazy_static;
extern crate rand;
extern crate regex;
extern crate rustc_serialize;

pub mod counts;
pub mod database;
pub mod schema;
pub mod models;
pub mod quotes;

use self::database::{establish_connection};
use self::quotes::{retrieve_quote, store_quotes, RetrievalRequest};
use self::counts::{count_authors, count_quotes};
use std::fs::File;
use std::io::{BufReader, BufRead};
use std::path::Path;

docopt!(Args derive Debug, "
Quote Storage & Retrieval Utilities

Usage:
  quote_storage count [authors | unretrieved-quotes | retrieved-quotes]
  quote_storage store <file> [--mark-retrieved]
  quote_storage retrieve <quote-id> [--mark-retrieved]
  quote_storage retrieve [--random-retrieved | --random-unretrieved | --first-unretrieved] [--mark-retrieved]
  quote_storage (-h | --help)
  quote_storage --version

Options:
  -h --help             Show this screen.
  --version             Show version.
  --random              Retrieves a random quote.
  --random-retrieved    Retrieves a random already retrieved quote.
  --random-unretrieved  Retrieves a random quote that has not already been retrieved.
  --mark-retrieved      Marks quotes as retrieved.
");

fn main() {
    let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());

    if args.flag_version {
        let version = env!("CARGO_PKG_VERSION");

        return println!("quote_storage v{}", version);
    }

    let dbsession = establish_connection();

    // dbsession.begin_test_transaction().unwrap(); // Testing

    if args.cmd_store {
        let path = Path::new(&args.arg_file);
        let file = match File::open(&path) {
            Ok(f) => f,
            Err(e) => panic!("FIXME: Invalid file reference found :'( {}", e)
        };

        let quotes = quotes_from_buffered_reader(BufReader::new(file));

        store_quotes(dbsession, &quotes, args.flag_mark_retrieved);

    } else if args.cmd_retrieve {
        let request = if args.flag_random_retrieved {
            RetrievalRequest::RandomRetrieved
        } else if args.flag_random_unretrieved {
            RetrievalRequest::RandomUnretrievedAndMark(args.flag_mark_retrieved)
        } else if args.flag_first_unretrieved {
            RetrievalRequest::FirstUnretrievedAndMark(args.flag_mark_retrieved)
        } else if args.arg_quote_id.len() > 0 {
            let id = args.arg_quote_id.parse::<i64>().unwrap(); // FIXME: Handle error gracefully

            RetrievalRequest::IdAndMark(id, args.flag_mark_retrieved)
        } else {
            RetrievalRequest::Random
        };

        let quote = retrieve_quote(dbsession, request);

        println!("{}", quote);
    } else if args.cmd_count {
        let count = if args.cmd_authors {
            count_authors(dbsession)
        } else if args.cmd_unretrieved_quotes {
            count_quotes(dbsession, true)
        } else {
            count_quotes(dbsession, false)
        };

        print!("{}", count);
    } else {
        unreachable!("Should not be able to get here!")
    }
}

fn quotes_from_buffered_reader<B: BufRead>(bufreader: B) -> Vec<String> {
    bufreader.lines().map(|l| l.expect("Could not parse line")).collect()
}
<|endoftext|>