<commit_before>//! This file is used by the build script. Therefore all functions generating the app command line parser must be included
//! here. It would be nicer to move at least the subcommands to the corresponding modules, but then all logic, all crates
//! etc. used there have to be available for the build script which makes it much more complex.

use std::str::FromStr;

use clap::{crate_name, App, AppSettings, Arg, SubCommand};

pub static SUPPORTED_RECORD_TYPES: &[&str] = &[
    "A", "AAAA", "ANAME", "ANY", "CNAME", "MX", "NULL", "NS", "PTR", "SOA", "SRV", "TXT",
];

pub static SUPPORTED_OUTPUT_FORMATS: &[&str] = &["json", "summary"];

pub fn create_parser() -> App<'static, 'static> {
    App::new(crate_name!())
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .global_setting(AppSettings::DeriveDisplayOrder)
        .global_setting(AppSettings::DisableHelpSubcommand)
        .global_setting(AppSettings::GlobalVersion)
        .global_setting(AppSettings::InferSubcommands)
        .global_setting(AppSettings::UnifiedHelpMessage)
        .arg(
            Arg::with_name("no-system-resolv-opt")
                .long("no-system-resolv-opt")
                .help("Ignores options set in /etc/resolv.conf"),
        )
        .arg(
            Arg::with_name("no-system-nameservers")
                .long("no-system-nameservers")
                .requires("system nameservers")
                .help("Ignores nameservers from /etc/resolv.conf"),
        )
        .arg(
            Arg::with_name("resolv-conf")
                .long("resolv-conf")
                .value_name("FILE")
                .takes_value(true)
                .help("Uses alternative resolv.conf file"),
        )
        .arg(
            Arg::with_name("system nameservers")
                .short("S")
                .long("system-nameserver")
                .value_name("IP ADDR")
                .takes_value(true)
                .multiple(true)
                .help("Adds system nameserver for system lookups; only IP addresses allowed"),
        )
        .arg(
            Arg::with_name("nameservers")
                .short("s")
                .long("nameserver")
                .value_name("HOSTNAME | IP ADDR")
                .takes_value(true)
                .multiple(true)
                .use_delimiter(true)
                .require_delimiter(true)
                .value_delimiter(";")
                .help("Adds nameserver for lookups"),
        )
        .arg(
            Arg::with_name("predefined")
                .short("p")
                .long("predefined")
                .help("Adds predefined nameservers for lookups"),
        )
        .arg(
            Arg::with_name("predefined-filter")
                .long("predefined-filter")
                .value_name("PROTOCOL")
                .multiple(true)
                .use_delimiter(true)
                .require_delimiter(true)
                .default_value("udp,tcp,https,tls")
                .possible_values(&["udp", "tcp", "https", "tls"])
                .default_value_if("predefined", None, "udp")
                .help("Filters predefined nameservers by protocol"),
        )
        .arg(
            Arg::with_name("list-predefined")
                .long("list-predefined")
                .help("Lists all predefined nameservers"),
        )
        .arg(
            Arg::with_name("nameservers-from-file")
                .short("f")
                .long("nameservers-from-file")
                .value_name("FILE")
                .takes_value(true)
                .help("Adds nameserver for lookups from file"),
        )
        .arg(
            Arg::with_name("limit")
                .long("limit")
                .value_name("NUMBER")
                .default_value("100")
                .validator(|str| {
                    usize::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets max. number of nameservers to query"),
        )
        .arg(
            Arg::with_name("max-concurrent-servers")
                .long("max-concurrent-servers")
                .value_name("NUMBER")
                .default_value("10")
                .validator(|str| {
                    usize::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets max. concurrent nameservers"),
        )
        .arg(
            Arg::with_name("max-concurrent-requests")
                .long("max-concurrent-requests")
                .value_name("NUMBER")
                .default_value("5")
                .validator(|str| {
                    usize::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets max. concurrent requests per nameserver"),
        )
        .arg(
            Arg::with_name("retries")
                .long("retries")
                .value_name("RETRIES")
                .default_value("0")
                .validator(|str| {
                    usize::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets number of retries if first lookup to nameserver fails"),
        )
        .arg(
            Arg::with_name("timeout")
                .long("timeout")
                .value_name("TIMEOUT")
                .default_value("5")
                .validator(|str| {
                    u64::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets timeout in seconds for responses"),
        )
        .arg(
            Arg::with_name("wait-multiple-responses")
                .long("wait-multiple-responses")
                .help("Waits until timeout for additional responses from nameservers"),
        )
        .arg(
            Arg::with_name("no-abort-on-error")
                .long("no-abort-on-error")
                .help("Sets do-not-ignore errors from nameservers"),
        )
        .arg(
            Arg::with_name("no-abort-on-timeout")
                .long("no-abort-on-timeout")
                .help("Sets do-not-ignore timeouts from nameservers"),
        )
        .arg(
            Arg::with_name("no-aborts")
                .long("no-aborts")
                .help("Sets do-not-ignore errors and timeouts from nameservers"),
        )
        .arg(
            Arg::with_name("output")
                .short("o")
                .long("output")
                .value_name("FORMAT")
                .takes_value(true)
                .default_value("summary")
                .possible_values(SUPPORTED_OUTPUT_FORMATS)
                .help("Sets the output format for result presentation"),
        )
        .arg(
            Arg::with_name("output-options")
                .long("output-options")
                .value_name("OPTIONS")
                .multiple(true)
                .use_delimiter(true)
                .require_delimiter(true)
                .default_value_if("output", Some("json"), "pretty")
                .default_value_if("output", Some("summary"), "human")
                .help("Sets output options")
                .long_help(
                    "* Json:
  * 'pretty': Prettifies output
* Summary:
  * 'condensed': Simplifies output,
  * 'human': Uses human readable formatting
  * 'show-domain-names': Shows queried domain names",
                ),
        )
        .arg(
            Arg::with_name("show-errors")
                .long("show-errors")
                .conflicts_with("quiet")
                .help("Shows error counts"),
        )
        .arg(
            Arg::with_name("quiet")
                .short("q")
                .long("quiet")
                .help("Does not print anything but results"),
        )
        // This is a special option that is not reflected in GlobalConfig, but is checked during
        // setup in `mhost.rs`.
        .arg(
            Arg::with_name("no-color")
                .long("no-color")
                .help("Disables colorful output"),
        )
        // This is a special option that is not reflected in GlobalConfig, but is checked during
        // setup in `mhost.rs` and set the global AtomicBool `mhost::output::styles::ASCII_MODE`.
        .arg(
            Arg::with_name("ascii")
                .long("ascii")
                .help("Uses only ASCII compatible characters for output"),
        )
        .arg(
            Arg::with_name("v")
                .short("v")
                .multiple(true)
                .help("Sets the level of verbosity"),
        )
        .subcommands(subcommands())
}

fn subcommands() -> Vec<App<'static, 'static>> {
    vec![
        discover_subcommand(),
        get_server_lists_subcommand(),
        lookup_subcommand(),
        soa_check_subcommand(),
    ]
    .into_iter()
    .map(|x| x.version(env!("CARGO_PKG_VERSION")).author(env!("CARGO_PKG_AUTHORS")))
    .collect()
}

fn discover_subcommand() -> App<'static, 'static> {
    SubCommand::with_name("discover")
        .about("Discovers records of a domain using multiple heuristics")
        .arg(
            Arg::with_name("domain name")
                .index(1)
                .value_name("NAME")
                .next_line_help(false)
                .help("domain name to check")
                .long_help("* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de"),
        )
        .arg(
            Arg::with_name("partial-results")
                .short("p")
                .long("show-partial-results")
                .help("Shows results after each lookup step"),
        )
        .arg(
            Arg::with_name("single-server-lookup")
                .short("S")
                .long("single-server-lookup")
                .help("Switches into single server lookup mode: every query will be send just one randomly chosen nameserver. This can be used to distribute queries among the available nameservers.")
        )
        .arg(
            Arg::with_name("wordlist-from-file")
                .short("w")
                .long("wordlist-from-file")
                .value_name("FILE")
                .takes_value(true)
                .help("Uses wordlist from file"),
        )
        .arg(
            Arg::with_name("rnd-names-number")
                .long("rnd-names-number")
                .value_name("NUMBER")
                .default_value("3")
                .validator(|str| {
                    u64::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets number of random domain names to generate for wildcard resolution check")
        )
        .arg(
            Arg::with_name("rnd-names-len")
                .long("rnd-names-len")
                .value_name("len")
                .default_value("32")
                .validator(|str| {
                    u64::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets length of random domain names to generate for wildcard resolution check")
        )
}

fn get_server_lists_subcommand() -> App<'static, 'static> {
    SubCommand::with_name("get-server-lists")
        .about("Downloads known lists of name servers")
        .arg(
            Arg::with_name("server_list_spec")
                .index(1)
                .value_name("SERVER_LIST_SPEC")
                .multiple(true)
                .required(true)
                .next_line_help(false)
                .help("server list specification")
                .long_help(
                    r#"SERVER LIST SPEC as <SOURCE>[:OPTIONS,...]
* 'public-dns' with options - cf. https://public-dns.info
   Example: public-dns:de
  '<top level country domain>': options select servers from that country
* 'opennic' with options; uses GeoIP to select servers - cf. https://www.opennic.org
   'anon' - only return servers with anonymized logs only; default is false
   'number=<1..>' - return up to 'number' servers; default is 10
   'reliability=<1..100> - only return server with reliability of 'reliability'% or more; default 95
   'ipv=<4|6|all> - return IPv4, IPv6, or both servers; default all
    Example: opennic:anon,number=10,ipv=4
"#,
                ),
        )
        .arg(
            Arg::with_name("output-file")
                .short("o")
                .long("output-file")
                .required(true)
                .value_name("FILE")
                .takes_value(true)
                .help("Sets path to output file"),
        )
}

fn lookup_subcommand() -> App<'static, 'static> {
    SubCommand::with_name("lookup")
        .about("Looks up a name, IP address or CIDR block")
        .arg(Arg::with_name("domain name")
            .required_unless("list-predefined")
            .index(1)
            .value_name("NAME | IP ADDR | CIDR BLOCK")
            .next_line_help(false)
            .help("domain name, IP address, or CIDR block to lookup")
            .long_help(
                "* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de
* IP ADDR may be any valid IPv4 or IPv4 address, e.g., 192.168.0.1
* CIDR BLOCK may be any valid IPv4 or IPv6 subnet in CIDR notation, e.g., 192.168.0.1/24
  all valid IP addresses of a CIDR block will be queried for a reverse lookup")
        )
        .arg(Arg::with_name("record types")
            .short("t")
            .long("record-type")
            .value_name("RECORD TYPE")
            .takes_value(true)
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value("A,AAAA,MX")
            .possible_values(SUPPORTED_RECORD_TYPES)
            .help("Sets record type to lookup, will be ignored in case of IP address lookup")
        )
        .arg(Arg::with_name("all-record-types")
            .long("all")
            .alias("xmas")
            .help("Enables lookups for all record types")
        )
        .arg(Arg::with_name("single-server-lookup")
            .short("S")
            .long("single-server-lookup")
            .help("Switches into single server lookup mode: every query will be send just one randomly chosen nameserver. This can be used to distribute queries among the available nameservers.")
        )
        .arg(Arg::with_name("whois")
            .short("w")
            .long("whois")
            .help("Retrieves Whois information about A, AAAA, and PTR records.")
        )
}

fn soa_check_subcommand() -> App<'static, 'static> {
    SubCommand::with_name("soa-check")
        .about("Checks SOA records of authoritative name servers for deviations")
        .arg(
            Arg::with_name("domain name")
                .index(1)
                .value_name("NAME")
                .next_line_help(false)
                .help("domain name to check")
                .long_help("* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de"),
        )
        .arg(
            Arg::with_name("partial-results")
                .short("p")
                .long("show-partial-results")
                .help("Shows results after each lookup step"),
        )
}

pub fn show_help() {
    let _ = create_parser().print_help();
    // Add line break after last line of help
    println!();
}
<commit_msg>Fix default in CLI for predefined-filter<commit_after>//! This file is used by the build script. Therefore all functions generating the app command line parser must be included
//! here. It would be nicer to move at least the subcommands to the corresponding modules, but then all logic, all crates
//! etc. used there have to be available for the build script which makes it much more complex.

use std::str::FromStr;

use clap::{crate_name, App, AppSettings, Arg, SubCommand};

pub static SUPPORTED_RECORD_TYPES: &[&str] = &[
    "A", "AAAA", "ANAME", "ANY", "CNAME", "MX", "NULL", "NS", "PTR", "SOA", "SRV", "TXT",
];

pub static SUPPORTED_OUTPUT_FORMATS: &[&str] = &["json", "summary"];

pub fn create_parser() -> App<'static, 'static> {
    App::new(crate_name!())
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .global_setting(AppSettings::DeriveDisplayOrder)
        .global_setting(AppSettings::DisableHelpSubcommand)
        .global_setting(AppSettings::GlobalVersion)
        .global_setting(AppSettings::InferSubcommands)
        .global_setting(AppSettings::UnifiedHelpMessage)
        .arg(
            Arg::with_name("no-system-resolv-opt")
                .long("no-system-resolv-opt")
                .help("Ignores options set in /etc/resolv.conf"),
        )
        .arg(
            Arg::with_name("no-system-nameservers")
                .long("no-system-nameservers")
                .requires("system nameservers")
                .help("Ignores nameservers from /etc/resolv.conf"),
        )
        .arg(
            Arg::with_name("resolv-conf")
                .long("resolv-conf")
                .value_name("FILE")
                .takes_value(true)
                .help("Uses alternative resolv.conf file"),
        )
        .arg(
            Arg::with_name("system nameservers")
                .short("S")
                .long("system-nameserver")
                .value_name("IP ADDR")
                .takes_value(true)
                .multiple(true)
                .help("Adds system nameserver for system lookups; only IP addresses allowed"),
        )
        .arg(
            Arg::with_name("nameservers")
                .short("s")
                .long("nameserver")
                .value_name("HOSTNAME | IP ADDR")
                .takes_value(true)
                .multiple(true)
                .use_delimiter(true)
                .require_delimiter(true)
                .value_delimiter(";")
                .help("Adds nameserver for lookups"),
        )
        .arg(
            Arg::with_name("predefined")
                .short("p")
                .long("predefined")
                .help("Adds predefined nameservers for lookups"),
        )
        .arg(
            Arg::with_name("predefined-filter")
                .long("predefined-filter")
                .value_name("PROTOCOL")
                .multiple(true)
                .use_delimiter(true)
                .require_delimiter(true)
                .default_value("udp")
                .possible_values(&["udp", "tcp", "https", "tls"])
                .default_value_if("predefined", None, "udp")
                .help("Filters predefined nameservers by protocol"),
        )
        .arg(
            Arg::with_name("list-predefined")
                .long("list-predefined")
                .help("Lists all predefined nameservers"),
        )
        .arg(
            Arg::with_name("nameservers-from-file")
                .short("f")
                .long("nameservers-from-file")
                .value_name("FILE")
                .takes_value(true)
                .help("Adds nameserver for lookups from file"),
        )
        .arg(
            Arg::with_name("limit")
                .long("limit")
                .value_name("NUMBER")
                .default_value("100")
                .validator(|str| {
                    usize::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets max. number of nameservers to query"),
        )
        .arg(
            Arg::with_name("max-concurrent-servers")
                .long("max-concurrent-servers")
                .value_name("NUMBER")
                .default_value("10")
                .validator(|str| {
                    usize::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets max. concurrent nameservers"),
        )
        .arg(
            Arg::with_name("max-concurrent-requests")
                .long("max-concurrent-requests")
                .value_name("NUMBER")
                .default_value("5")
                .validator(|str| {
                    usize::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets max. concurrent requests per nameserver"),
        )
        .arg(
            Arg::with_name("retries")
                .long("retries")
                .value_name("RETRIES")
                .default_value("0")
                .validator(|str| {
                    usize::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets number of retries if first lookup to nameserver fails"),
        )
        .arg(
            Arg::with_name("timeout")
                .long("timeout")
                .value_name("TIMEOUT")
                .default_value("5")
                .validator(|str| {
                    u64::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets timeout in seconds for responses"),
        )
        .arg(
            Arg::with_name("wait-multiple-responses")
                .long("wait-multiple-responses")
                .help("Waits until timeout for additional responses from nameservers"),
        )
        .arg(
            Arg::with_name("no-abort-on-error")
                .long("no-abort-on-error")
                .help("Sets do-not-ignore errors from nameservers"),
        )
        .arg(
            Arg::with_name("no-abort-on-timeout")
                .long("no-abort-on-timeout")
                .help("Sets do-not-ignore timeouts from nameservers"),
        )
        .arg(
            Arg::with_name("no-aborts")
                .long("no-aborts")
                .help("Sets do-not-ignore errors and timeouts from nameservers"),
        )
        .arg(
            Arg::with_name("output")
                .short("o")
                .long("output")
                .value_name("FORMAT")
                .takes_value(true)
                .default_value("summary")
                .possible_values(SUPPORTED_OUTPUT_FORMATS)
                .help("Sets the output format for result presentation"),
        )
        .arg(
            Arg::with_name("output-options")
                .long("output-options")
                .value_name("OPTIONS")
                .multiple(true)
                .use_delimiter(true)
                .require_delimiter(true)
                .default_value_if("output", Some("json"), "pretty")
                .default_value_if("output", Some("summary"), "human")
                .help("Sets output options")
                .long_help(
                    "* Json:
  * 'pretty': Prettifies output
* Summary:
  * 'condensed': Simplifies output,
  * 'human': Uses human readable formatting
  * 'show-domain-names': Shows queried domain names",
                ),
        )
        .arg(
            Arg::with_name("show-errors")
                .long("show-errors")
                .conflicts_with("quiet")
                .help("Shows error counts"),
        )
        .arg(
            Arg::with_name("quiet")
                .short("q")
                .long("quiet")
                .help("Does not print anything but results"),
        )
        // This is a special option that is not reflected in GlobalConfig, but is checked during
        // setup in `mhost.rs`.
        .arg(
            Arg::with_name("no-color")
                .long("no-color")
                .help("Disables colorful output"),
        )
        // This is a special option that is not reflected in GlobalConfig, but is checked during
        // setup in `mhost.rs` and set the global AtomicBool `mhost::output::styles::ASCII_MODE`.
        .arg(
            Arg::with_name("ascii")
                .long("ascii")
                .help("Uses only ASCII compatible characters for output"),
        )
        .arg(
            Arg::with_name("v")
                .short("v")
                .multiple(true)
                .help("Sets the level of verbosity"),
        )
        .subcommands(subcommands())
}

fn subcommands() -> Vec<App<'static, 'static>> {
    vec![
        discover_subcommand(),
        get_server_lists_subcommand(),
        lookup_subcommand(),
        soa_check_subcommand(),
    ]
    .into_iter()
    .map(|x| x.version(env!("CARGO_PKG_VERSION")).author(env!("CARGO_PKG_AUTHORS")))
    .collect()
}

fn discover_subcommand() -> App<'static, 'static> {
    SubCommand::with_name("discover")
        .about("Discovers records of a domain using multiple heuristics")
        .arg(
            Arg::with_name("domain name")
                .index(1)
                .value_name("NAME")
                .next_line_help(false)
                .help("domain name to check")
                .long_help("* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de"),
        )
        .arg(
            Arg::with_name("partial-results")
                .short("p")
                .long("show-partial-results")
                .help("Shows results after each lookup step"),
        )
        .arg(
            Arg::with_name("single-server-lookup")
                .short("S")
                .long("single-server-lookup")
                .help("Switches into single server lookup mode: every query will be send just one randomly chosen nameserver. This can be used to distribute queries among the available nameservers.")
        )
        .arg(
            Arg::with_name("wordlist-from-file")
                .short("w")
                .long("wordlist-from-file")
                .value_name("FILE")
                .takes_value(true)
                .help("Uses wordlist from file"),
        )
        .arg(
            Arg::with_name("rnd-names-number")
                .long("rnd-names-number")
                .value_name("NUMBER")
                .default_value("3")
                .validator(|str| {
                    u64::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets number of random domain names to generate for wildcard resolution check")
        )
        .arg(
            Arg::with_name("rnd-names-len")
                .long("rnd-names-len")
                .value_name("len")
                .default_value("32")
                .validator(|str| {
                    u64::from_str(&str)
                        .map(|_| ())
                        .map_err(|_| "invalid number".to_string())
                })
                .help("Sets length of random domain names to generate for wildcard resolution check")
        )
}

fn get_server_lists_subcommand() -> App<'static, 'static> {
    SubCommand::with_name("get-server-lists")
        .about("Downloads known lists of name servers")
        .arg(
            Arg::with_name("server_list_spec")
                .index(1)
                .value_name("SERVER_LIST_SPEC")
                .multiple(true)
                .required(true)
                .next_line_help(false)
                .help("server list specification")
                .long_help(
                    r#"SERVER LIST SPEC as <SOURCE>[:OPTIONS,...]
* 'public-dns' with options - cf. https://public-dns.info
   Example: public-dns:de
  '<top level country domain>': options select servers from that country
* 'opennic' with options; uses GeoIP to select servers - cf. https://www.opennic.org
   'anon' - only return servers with anonymized logs only; default is false
   'number=<1..>' - return up to 'number' servers; default is 10
   'reliability=<1..100> - only return server with reliability of 'reliability'% or more; default 95
   'ipv=<4|6|all> - return IPv4, IPv6, or both servers; default all
    Example: opennic:anon,number=10,ipv=4
"#,
                ),
        )
        .arg(
            Arg::with_name("output-file")
                .short("o")
                .long("output-file")
                .required(true)
                .value_name("FILE")
                .takes_value(true)
                .help("Sets path to output file"),
        )
}

fn lookup_subcommand() -> App<'static, 'static> {
    SubCommand::with_name("lookup")
        .about("Looks up a name, IP address or CIDR block")
        .arg(Arg::with_name("domain name")
            .required_unless("list-predefined")
            .index(1)
            .value_name("NAME | IP ADDR | CIDR BLOCK")
            .next_line_help(false)
            .help("domain name, IP address, or CIDR block to lookup")
            .long_help(
                "* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de
* IP ADDR may be any valid IPv4 or IPv4 address, e.g., 192.168.0.1
* CIDR BLOCK may be any valid IPv4 or IPv6 subnet in CIDR notation, e.g., 192.168.0.1/24
  all valid IP addresses of a CIDR block will be queried for a reverse lookup")
        )
        .arg(Arg::with_name("record types")
            .short("t")
            .long("record-type")
            .value_name("RECORD TYPE")
            .takes_value(true)
            .multiple(true)
            .use_delimiter(true)
            .require_delimiter(true)
            .default_value("A,AAAA,MX")
            .possible_values(SUPPORTED_RECORD_TYPES)
            .help("Sets record type to lookup, will be ignored in case of IP address lookup")
        )
        .arg(Arg::with_name("all-record-types")
            .long("all")
            .alias("xmas")
            .help("Enables lookups for all record types")
        )
        .arg(Arg::with_name("single-server-lookup")
            .short("S")
            .long("single-server-lookup")
            .help("Switches into single server lookup mode: every query will be send just one randomly chosen nameserver. This can be used to distribute queries among the available nameservers.")
        )
        .arg(Arg::with_name("whois")
            .short("w")
            .long("whois")
            .help("Retrieves Whois information about A, AAAA, and PTR records.")
        )
}

fn soa_check_subcommand() -> App<'static, 'static> {
    SubCommand::with_name("soa-check")
        .about("Checks SOA records of authoritative name servers for deviations")
        .arg(
            Arg::with_name("domain name")
                .index(1)
                .value_name("NAME")
                .next_line_help(false)
                .help("domain name to check")
                .long_help("* DOMAIN NAME may be any valid DNS name, e.g., lukas.pustina.de"),
        )
        .arg(
            Arg::with_name("partial-results")
                .short("p")
                .long("show-partial-results")
                .help("Shows results after each lookup step"),
        )
}

pub fn show_help() {
    let _ = create_parser().print_help();
    // Add line break after last line of help
    println!();
}
<|endoftext|>


<commit_before>use std::time::Instant;

/// A metrics counter storing only latest `N` records.
pub struct MetricsCounter<const N: usize> {
    /// Slots to store metrics.
    slots: [(usize, Instant); N],
    /// The slot of the oldest record.
    /// Also the next slot to store the new record.
    index: usize,
}

impl<const N: usize> MetricsCounter<N> {
    /// Creates a new counter with an initial value.
    pub fn new(init: usize, init_at: Instant) -> Self {
        debug_assert!(N > 0, "number of slots must be greater than zero");
        Self {
            slots: [(init, init_at); N],
            index: 0,
        }
    }

    /// Adds record to the counter.
    pub fn add(&mut self, data: usize, added_at: Instant) {
        self.slots[self.index] = (data, added_at);
        self.index = (self.index + 1) % N;
    }

    /// Calculates per-second average rate of all slots.
    pub fn rate(&self) -> f32 {
        let latest = self.slots[self.index.checked_sub(1).unwrap_or(N - 1)];
        let oldest = self.slots[self.index];
        let duration = (latest.1 - oldest.1).as_secs_f32();
        let avg = (latest.0 - oldest.0) as f32 / duration;
        if f32::is_nan(avg) {
            0f32
        } else {
            avg
        }
    }
}

#[cfg(test)]
mod tests {
    use super::MetricsCounter;
    use std::time::{Duration, Instant};

    #[test]
    fn counter() {
        let now = Instant::now();
        let mut counter = MetricsCounter::<3>::new(0, now);
        assert_eq!(counter.rate(), 0f32);
        counter.add(1, now + Duration::from_secs(1));
        assert_eq!(counter.rate(), 1f32);
        counter.add(4, now + Duration::from_secs(2));
        assert_eq!(counter.rate(), 2f32);
        counter.add(7, now + Duration::from_secs(3));
        assert_eq!(counter.rate(), 3f32);
        counter.add(12, now + Duration::from_secs(4));
        assert_eq!(counter.rate(), 4f32);
    }

    #[test]
    #[should_panic(expected = "number of slots must be greater than zero")]
    fn counter_zero_slot() {
        let _counter = MetricsCounter::<0>::new(0, Instant::now());
    }
}
<commit_msg>Auto merge of #9440 - ehuss:fix-metrics-test, r=Eh2406<commit_after>use std::time::Instant;

/// A metrics counter storing only latest `N` records.
pub struct MetricsCounter<const N: usize> {
    /// Slots to store metrics.
    slots: [(usize, Instant); N],
    /// The slot of the oldest record.
    /// Also the next slot to store the new record.
    index: usize,
}

impl<const N: usize> MetricsCounter<N> {
    /// Creates a new counter with an initial value.
    pub fn new(init: usize, init_at: Instant) -> Self {
        assert!(N > 0, "number of slots must be greater than zero");
        Self {
            slots: [(init, init_at); N],
            index: 0,
        }
    }

    /// Adds record to the counter.
    pub fn add(&mut self, data: usize, added_at: Instant) {
        self.slots[self.index] = (data, added_at);
        self.index = (self.index + 1) % N;
    }

    /// Calculates per-second average rate of all slots.
    pub fn rate(&self) -> f32 {
        let latest = self.slots[self.index.checked_sub(1).unwrap_or(N - 1)];
        let oldest = self.slots[self.index];
        let duration = (latest.1 - oldest.1).as_secs_f32();
        let avg = (latest.0 - oldest.0) as f32 / duration;
        if f32::is_nan(avg) {
            0f32
        } else {
            avg
        }
    }
}

#[cfg(test)]
mod tests {
    use super::MetricsCounter;
    use std::time::{Duration, Instant};

    #[test]
    fn counter() {
        let now = Instant::now();
        let mut counter = MetricsCounter::<3>::new(0, now);
        assert_eq!(counter.rate(), 0f32);
        counter.add(1, now + Duration::from_secs(1));
        assert_eq!(counter.rate(), 1f32);
        counter.add(4, now + Duration::from_secs(2));
        assert_eq!(counter.rate(), 2f32);
        counter.add(7, now + Duration::from_secs(3));
        assert_eq!(counter.rate(), 3f32);
        counter.add(12, now + Duration::from_secs(4));
        assert_eq!(counter.rate(), 4f32);
    }

    #[test]
    #[should_panic(expected = "number of slots must be greater than zero")]
    fn counter_zero_slot() {
        let _counter = MetricsCounter::<0>::new(0, Instant::now());
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::path::Path;

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let out_path = Path::new(&out_dir).join("bindgen.rs");
    build::main(&out_dir, &out_path);
}

#[cfg(feature = "bundled")]
mod build {
    use cc;
    use std::path::Path;

    pub fn main(out_dir: &str, out_path: &Path) {
        if cfg!(feature = "sqlcipher") {
            panic!("Builds with bundled SQLCipher are not supported");
        }

        #[cfg(feature = "buildtime_bindgen")]
        {
            use super::{bindings, HeaderLocation};
            let header = HeaderLocation::FromPath("sqlite3/sqlite3.h".to_owned());
            bindings::write_to_out_dir(header, out_path);
        }
        #[cfg(not(feature = "buildtime_bindgen"))]
        {
            use std::fs;
            fs::copy("sqlite3/bindgen_bundled_version.rs", out_path)
                .expect("Could not copy bindings to output directory");
        }

        let mut cfg = cc::Build::new();
        cfg.file("sqlite3/sqlite3.c")
            .flag("-DSQLITE_CORE")
            .flag("-DSQLITE_DEFAULT_FOREIGN_KEYS=1")
            .flag("-DSQLITE_ENABLE_API_ARMOR")
            .flag("-DSQLITE_ENABLE_COLUMN_METADATA")
            .flag("-DSQLITE_ENABLE_DBSTAT_VTAB")
            .flag("-DSQLITE_ENABLE_FTS3")
            .flag("-DSQLITE_ENABLE_FTS3_PARENTHESIS")
            .flag("-DSQLITE_ENABLE_FTS5")
            .flag("-DSQLITE_ENABLE_JSON1")
            .flag("-DSQLITE_ENABLE_LOAD_EXTENSION=1")
            .flag("-DSQLITE_ENABLE_MEMORY_MANAGEMENT")
            .flag("-DSQLITE_ENABLE_RTREE")
            .flag("-DSQLITE_ENABLE_STAT2")
            .flag("-DSQLITE_ENABLE_STAT4")
            .flag("-DSQLITE_HAVE_ISNAN")
            .flag("-DSQLITE_SOUNDEX")
            .flag("-DSQLITE_THREADSAFE=1")
            .flag("-DSQLITE_USE_URI")
            .flag("-DHAVE_USLEEP=1");
        if cfg!(feature = "unlock_notify") {
            cfg.flag("-DSQLITE_ENABLE_UNLOCK_NOTIFY");
        }
        if cfg!(feature = "preupdate_hook") {
            cfg.flag("-DSQLITE_ENABLE_PREUPDATE_HOOK");
        }
        if cfg!(feature = "session") {
            cfg.flag("-DSQLITE_ENABLE_SESSION");
        }
        cfg.compile("libsqlite3.a");

        println!("cargo:lib_dir={}", out_dir);
    }
}

fn env_prefix() -> &'static str {
    if cfg!(feature = "sqlcipher") {
        "SQLCIPHER"
    } else {
        "SQLITE3"
    }
}

pub enum HeaderLocation {
    FromEnvironment,
    Wrapper,
    FromPath(String),
}

impl From<HeaderLocation> for String {
    fn from(header: HeaderLocation) -> String {
        match header {
            HeaderLocation::FromEnvironment => {
                let prefix = env_prefix();
                let mut header = env::var(format!("{}_INCLUDE_DIR", prefix)).expect(&format!(
                    "{}_INCLUDE_DIR must be set if {}_LIB_DIR is set",
                    prefix, prefix
                ));
                header.push_str("/sqlite3.h");
                header
            }
            HeaderLocation::Wrapper => "wrapper.h".into(),
            HeaderLocation::FromPath(path) => path,
        }
    }
}

#[cfg(not(feature = "bundled"))]
mod build {
    use pkg_config;

    #[cfg(all(feature = "vcpkg", target_env = "msvc"))]
    extern crate vcpkg;

    use super::{bindings, env_prefix, HeaderLocation};
    use std::env;
    use std::path::Path;

    pub fn main(_out_dir: &str, out_path: &Path) {
        let header = find_sqlite();
        bindings::write_to_out_dir(header, out_path);
    }

    fn find_link_mode() -> &'static str {
        // If the user specifies SQLITE_STATIC (or SQLCIPHER_STATIC), do static
        // linking, unless it's explicitly set to 0.
        match &env::var(format!("{}_STATIC", env_prefix())) {
            Ok(v) if v != "0" => "static",
            _ => "dylib",
        }
    }
    // Prints the necessary cargo link commands and returns the path to the header.
    fn find_sqlite() -> HeaderLocation {
        let link_lib = link_lib();

        println!("cargo:rerun-if-env-changed={}_INCLUDE_DIR", env_prefix());
        println!("cargo:rerun-if-env-changed={}_LIB_DIR", env_prefix());
        println!("cargo:rerun-if-env-changed={}_STATIC", env_prefix());
        if cfg!(target_os = "windows") {
            println!("cargo:rerun-if-env-changed=PATH");
        }
        // Allow users to specify where to find SQLite.
        if let Ok(dir) = env::var(format!("{}_LIB_DIR", env_prefix())) {
            println!("cargo:rustc-link-lib={}={}", find_link_mode(), link_lib);
            println!("cargo:rustc-link-search={}", dir);
            return HeaderLocation::FromEnvironment;
        }

        if let Some(header) = try_vcpkg() {
            return header;
        }

        // See if pkg-config can do everything for us.
        match pkg_config::Config::new()
            .print_system_libs(false)
            .probe(link_lib)
        {
            Ok(mut lib) => {
                if let Some(mut header) = lib.include_paths.pop() {
                    header.push("sqlite3.h");
                    HeaderLocation::FromPath(header.to_string_lossy().into())
                } else {
                    HeaderLocation::Wrapper
                }
            }
            Err(_) => {
                // No env var set and pkg-config couldn't help; just output the link-lib
                // request and hope that the library exists on the system paths. We used to
                // output /usr/lib explicitly, but that can introduce other linking problems;
                // see https://github.com/jgallagher/rusqlite/issues/207.
                println!("cargo:rustc-link-lib={}={}", find_link_mode(), link_lib);
                HeaderLocation::Wrapper
            }
        }
    }

    #[cfg(all(feature = "vcpkg", target_env = "msvc"))]
    fn try_vcpkg() -> Option<HeaderLocation> {
        // See if vcpkg can find it.
        if let Ok(mut lib) = vcpkg::Config::new().probe(link_lib()) {
            if let Some(mut header) = lib.include_paths.pop() {
                header.push("sqlite3.h");
                return Some(HeaderLocation::FromPath(header.to_string_lossy().into()));
            }
        }
        None
    }

    #[cfg(not(all(feature = "vcpkg", target_env = "msvc")))]
    fn try_vcpkg() -> Option<HeaderLocation> {
        None
    }

    fn link_lib() -> &'static str {
        if cfg!(feature = "sqlcipher") {
            "sqlcipher"
        } else {
            "sqlite3"
        }
    }
}

#[cfg(all(not(feature = "buildtime_bindgen"), not(feature = "bundled")))]
mod bindings {
    use super::HeaderLocation;

    use std::fs;
    use std::path::Path;

    static PREBUILT_BINDGEN_PATHS: &'static [&'static str] = &[
        "bindgen-bindings/bindgen_3.6.8.rs",
        #[cfg(feature = "min_sqlite_version_3_6_23")]
        "bindgen-bindings/bindgen_3.6.23.rs",
        #[cfg(feature = "min_sqlite_version_3_7_7")]
        "bindgen-bindings/bindgen_3.7.7.rs",
        #[cfg(feature = "min_sqlite_version_3_7_16")]
        "bindgen-bindings/bindgen_3.7.16.rs",
    ];

    pub fn write_to_out_dir(_header: HeaderLocation, out_path: &Path) {
        let in_path = PREBUILT_BINDGEN_PATHS[PREBUILT_BINDGEN_PATHS.len() - 1];
        fs::copy(in_path, out_path).expect("Could not copy bindings to output directory");
    }
}

#[cfg(feature = "buildtime_bindgen")]
mod bindings {
    use bindgen;

    use self::bindgen::callbacks::{IntKind, ParseCallbacks};
    use super::HeaderLocation;

    use std::fs::OpenOptions;
    use std::io::Write;
    use std::path::Path;

    #[derive(Debug)]
    struct SqliteTypeChooser;

    impl ParseCallbacks for SqliteTypeChooser {
        fn int_macro(&self, _name: &str, value: i64) -> Option<IntKind> {
            if value >= i32::min_value() as i64 && value <= i32::max_value() as i64 {
                Some(IntKind::I32)
            } else {
                None
            }
        }
    }

    pub fn write_to_out_dir(header: HeaderLocation, out_path: &Path) {
        let header: String = header.into();
        let mut output = Vec::new();
        let mut bindings = bindgen::builder()
            .header(header.clone())
            .parse_callbacks(Box::new(SqliteTypeChooser))
            .rustfmt_bindings(true);

        if cfg!(feature = "unlock_notify") {
            bindings = bindings.clang_arg("-DSQLITE_ENABLE_UNLOCK_NOTIFY");
        }
        if cfg!(feature = "preupdate_hook") {
            bindings = bindings.clang_arg("-DSQLITE_ENABLE_PREUPDATE_HOOK");
        }
        if cfg!(feature = "session") {
            bindings = bindings.clang_arg("-DSQLITE_ENABLE_SESSION");
        }

        bindings
            .generate()
            .expect(&format!("could not run bindgen on header {}", header))
            .write(Box::new(&mut output))
            .expect("could not write output of bindgen");
        let mut output = String::from_utf8(output).expect("bindgen output was not UTF-8?!");

        // rusqlite's functions feature ors in the SQLITE_DETERMINISTIC flag when it
        // can. This flag was added in SQLite 3.8.3, but oring it in in prior
        // versions of SQLite is harmless. We don't want to not build just
        // because this flag is missing (e.g., if we're linking against
        // SQLite 3.7.x), so append the flag manually if it isn't present in bindgen's
        // output.
        if !output.contains("pub const SQLITE_DETERMINISTIC") {
            output.push_str("\npub const SQLITE_DETERMINISTIC: i32 = 2048;\n");
        }

        let mut file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(out_path.clone())
            .expect(&format!("Could not write to {:?}", out_path));

        file.write_all(output.as_bytes())
            .expect(&format!("Could not write to {:?}", out_path));
    }
}
<commit_msg>Use pkg-config to generate sqlite link dependencies<commit_after>use std::env;
use std::path::Path;

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let out_path = Path::new(&out_dir).join("bindgen.rs");
    build::main(&out_dir, &out_path);
}

#[cfg(feature = "bundled")]
mod build {
    use cc;
    use std::path::Path;

    pub fn main(out_dir: &str, out_path: &Path) {
        if cfg!(feature = "sqlcipher") {
            panic!("Builds with bundled SQLCipher are not supported");
        }

        #[cfg(feature = "buildtime_bindgen")]
        {
            use super::{bindings, HeaderLocation};
            let header = HeaderLocation::FromPath("sqlite3/sqlite3.h".to_owned());
            bindings::write_to_out_dir(header, out_path);
        }
        #[cfg(not(feature = "buildtime_bindgen"))]
        {
            use std::fs;
            fs::copy("sqlite3/bindgen_bundled_version.rs", out_path)
                .expect("Could not copy bindings to output directory");
        }

        let mut cfg = cc::Build::new();
        cfg.file("sqlite3/sqlite3.c")
            .flag("-DSQLITE_CORE")
            .flag("-DSQLITE_DEFAULT_FOREIGN_KEYS=1")
            .flag("-DSQLITE_ENABLE_API_ARMOR")
            .flag("-DSQLITE_ENABLE_COLUMN_METADATA")
            .flag("-DSQLITE_ENABLE_DBSTAT_VTAB")
            .flag("-DSQLITE_ENABLE_FTS3")
            .flag("-DSQLITE_ENABLE_FTS3_PARENTHESIS")
            .flag("-DSQLITE_ENABLE_FTS5")
            .flag("-DSQLITE_ENABLE_JSON1")
            .flag("-DSQLITE_ENABLE_LOAD_EXTENSION=1")
            .flag("-DSQLITE_ENABLE_MEMORY_MANAGEMENT")
            .flag("-DSQLITE_ENABLE_RTREE")
            .flag("-DSQLITE_ENABLE_STAT2")
            .flag("-DSQLITE_ENABLE_STAT4")
            .flag("-DSQLITE_HAVE_ISNAN")
            .flag("-DSQLITE_SOUNDEX")
            .flag("-DSQLITE_THREADSAFE=1")
            .flag("-DSQLITE_USE_URI")
            .flag("-DHAVE_USLEEP=1");
        if cfg!(feature = "unlock_notify") {
            cfg.flag("-DSQLITE_ENABLE_UNLOCK_NOTIFY");
        }
        if cfg!(feature = "preupdate_hook") {
            cfg.flag("-DSQLITE_ENABLE_PREUPDATE_HOOK");
        }
        if cfg!(feature = "session") {
            cfg.flag("-DSQLITE_ENABLE_SESSION");
        }
        cfg.compile("libsqlite3.a");

        println!("cargo:lib_dir={}", out_dir);
    }
}

fn env_prefix() -> &'static str {
    if cfg!(feature = "sqlcipher") {
        "SQLCIPHER"
    } else {
        "SQLITE3"
    }
}

pub enum HeaderLocation {
    FromEnvironment,
    Wrapper,
    FromPath(String),
}

impl From<HeaderLocation> for String {
    fn from(header: HeaderLocation) -> String {
        match header {
            HeaderLocation::FromEnvironment => {
                let prefix = env_prefix();
                let mut header = env::var(format!("{}_INCLUDE_DIR", prefix)).expect(&format!(
                    "{}_INCLUDE_DIR must be set if {}_LIB_DIR is set",
                    prefix, prefix
                ));
                header.push_str("/sqlite3.h");
                header
            }
            HeaderLocation::Wrapper => "wrapper.h".into(),
            HeaderLocation::FromPath(path) => path,
        }
    }
}

#[cfg(not(feature = "bundled"))]
mod build {
    use pkg_config;

    #[cfg(all(feature = "vcpkg", target_env = "msvc"))]
    extern crate vcpkg;

    use super::{bindings, env_prefix, HeaderLocation};
    use std::env;
    use std::path::Path;

    pub fn main(_out_dir: &str, out_path: &Path) {
        let header = find_sqlite();
        bindings::write_to_out_dir(header, out_path);
    }

    fn find_link_mode() -> &'static str {
        // If the user specifies SQLITE_STATIC (or SQLCIPHER_STATIC), do static
        // linking, unless it's explicitly set to 0.
        match &env::var(format!("{}_STATIC", env_prefix())) {
            Ok(v) if v != "0" => "static",
            _ => "dylib",
        }
    }
    // Prints the necessary cargo link commands and returns the path to the header.
    fn find_sqlite() -> HeaderLocation {
        let link_lib = link_lib();

        println!("cargo:rerun-if-env-changed={}_INCLUDE_DIR", env_prefix());
        println!("cargo:rerun-if-env-changed={}_LIB_DIR", env_prefix());
        println!("cargo:rerun-if-env-changed={}_STATIC", env_prefix());
        if cfg!(target_os = "windows") {
            println!("cargo:rerun-if-env-changed=PATH");
        }
        // Allow users to specify where to find SQLite.
        if let Ok(dir) = env::var(format!("{}_LIB_DIR", env_prefix())) {
            // Try to use pkg-config to determine link commands
            let pkgconfig_path = Path::new(&dir).join("pkgconfig");
            env::set_var("PKG_CONFIG_PATH", pkgconfig_path);
            if let Err(_) = pkg_config::Config::new().probe(link_lib) {
                // Otherwise just emit the bare minimum link commands.
                println!("cargo:rustc-link-lib={}={}", find_link_mode(), link_lib);
                println!("cargo:rustc-link-search={}", dir);
            }
            return HeaderLocation::FromEnvironment;
        }

        if let Some(header) = try_vcpkg() {
            return header;
        }

        // See if pkg-config can do everything for us.
        match pkg_config::Config::new()
            .print_system_libs(false)
            .probe(link_lib)
        {
            Ok(mut lib) => {
                if let Some(mut header) = lib.include_paths.pop() {
                    header.push("sqlite3.h");
                    HeaderLocation::FromPath(header.to_string_lossy().into())
                } else {
                    HeaderLocation::Wrapper
                }
            }
            Err(_) => {
                // No env var set and pkg-config couldn't help; just output the link-lib
                // request and hope that the library exists on the system paths. We used to
                // output /usr/lib explicitly, but that can introduce other linking problems;
                // see https://github.com/jgallagher/rusqlite/issues/207.
                println!("cargo:rustc-link-lib={}={}", find_link_mode(), link_lib);
                HeaderLocation::Wrapper
            }
        }
    }

    #[cfg(all(feature = "vcpkg", target_env = "msvc"))]
    fn try_vcpkg() -> Option<HeaderLocation> {
        // See if vcpkg can find it.
        if let Ok(mut lib) = vcpkg::Config::new().probe(link_lib()) {
            if let Some(mut header) = lib.include_paths.pop() {
                header.push("sqlite3.h");
                return Some(HeaderLocation::FromPath(header.to_string_lossy().into()));
            }
        }
        None
    }

    #[cfg(not(all(feature = "vcpkg", target_env = "msvc")))]
    fn try_vcpkg() -> Option<HeaderLocation> {
        None
    }

    fn link_lib() -> &'static str {
        if cfg!(feature = "sqlcipher") {
            "sqlcipher"
        } else {
            "sqlite3"
        }
    }
}

#[cfg(all(not(feature = "buildtime_bindgen"), not(feature = "bundled")))]
mod bindings {
    use super::HeaderLocation;

    use std::fs;
    use std::path::Path;

    static PREBUILT_BINDGEN_PATHS: &'static [&'static str] = &[
        "bindgen-bindings/bindgen_3.6.8.rs",
        #[cfg(feature = "min_sqlite_version_3_6_23")]
        "bindgen-bindings/bindgen_3.6.23.rs",
        #[cfg(feature = "min_sqlite_version_3_7_7")]
        "bindgen-bindings/bindgen_3.7.7.rs",
        #[cfg(feature = "min_sqlite_version_3_7_16")]
        "bindgen-bindings/bindgen_3.7.16.rs",
    ];

    pub fn write_to_out_dir(_header: HeaderLocation, out_path: &Path) {
        let in_path = PREBUILT_BINDGEN_PATHS[PREBUILT_BINDGEN_PATHS.len() - 1];
        fs::copy(in_path, out_path).expect("Could not copy bindings to output directory");
    }
}

#[cfg(feature = "buildtime_bindgen")]
mod bindings {
    use bindgen;

    use self::bindgen::callbacks::{IntKind, ParseCallbacks};
    use super::HeaderLocation;

    use std::fs::OpenOptions;
    use std::io::Write;
    use std::path::Path;

    #[derive(Debug)]
    struct SqliteTypeChooser;

    impl ParseCallbacks for SqliteTypeChooser {
        fn int_macro(&self, _name: &str, value: i64) -> Option<IntKind> {
            if value >= i32::min_value() as i64 && value <= i32::max_value() as i64 {
                Some(IntKind::I32)
            } else {
                None
            }
        }
    }

    pub fn write_to_out_dir(header: HeaderLocation, out_path: &Path) {
        let header: String = header.into();
        let mut output = Vec::new();
        let mut bindings = bindgen::builder()
            .header(header.clone())
            .parse_callbacks(Box::new(SqliteTypeChooser))
            .rustfmt_bindings(true);

        if cfg!(feature = "unlock_notify") {
            bindings = bindings.clang_arg("-DSQLITE_ENABLE_UNLOCK_NOTIFY");
        }
        if cfg!(feature = "preupdate_hook") {
            bindings = bindings.clang_arg("-DSQLITE_ENABLE_PREUPDATE_HOOK");
        }
        if cfg!(feature = "session") {
            bindings = bindings.clang_arg("-DSQLITE_ENABLE_SESSION");
        }

        bindings
            .generate()
            .expect(&format!("could not run bindgen on header {}", header))
            .write(Box::new(&mut output))
            .expect("could not write output of bindgen");
        let mut output = String::from_utf8(output).expect("bindgen output was not UTF-8?!");

        // rusqlite's functions feature ors in the SQLITE_DETERMINISTIC flag when it
        // can. This flag was added in SQLite 3.8.3, but oring it in in prior
        // versions of SQLite is harmless. We don't want to not build just
        // because this flag is missing (e.g., if we're linking against
        // SQLite 3.7.x), so append the flag manually if it isn't present in bindgen's
        // output.
        if !output.contains("pub const SQLITE_DETERMINISTIC") {
            output.push_str("\npub const SQLITE_DETERMINISTIC: i32 = 2048;\n");
        }

        let mut file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(out_path.clone())
            .expect(&format!("Could not write to {:?}", out_path));

        file.write_all(output.as_bytes())
            .expect(&format!("Could not write to {:?}", out_path));
    }
}
<|endoftext|>


<commit_before>

// Re-export
pub use uil_shared::markup::Node;
pub use uil_shared::markup::NodeType;
pub use uil_shared::markup::{Template, View};
pub use uil_shared::markup::{
    ButtonData,
    LineInputData,
    ProgressBarData,
    TemplateData,
    RepeatData
};

pub use uil_parsers::markup::Library;
pub use uil_parsers::markup::MAIN_VIEW_NAME;
<commit_msg>Missing re-export to build the example.<commit_after>

// Re-export
pub use uil_shared::markup::Node;
pub use uil_shared::markup::NodeType;
pub use uil_shared::markup::{Template, View};
pub use uil_shared::markup::{
    ButtonData,
    LineInputData,
    ProgressBarData,
    TemplateData,
    RepeatData
};

pub use uil_parsers::markup::Library;
pub use uil_parsers::markup::MAIN_VIEW_NAME;
pub use uil_parsers::markup::parse;
<|endoftext|>


<commit_before>//! Provide access to locally cached memdb SDKs
//!
//! The `MemDbStash` pulls in remote SDKs from an S3 bucket and provides
//! access to it.  This is used by the symbol server to manage the local
//! cache and also to refer to memdb files that are mmap'ed in.
use std::fs;
use std::io;
use std::iter::FromIterator;
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Values as HashMapValuesIter;
use std::sync::{Arc, RwLock};

use serde_json;
use xz2::write::XzDecoder;
use chrono::UTC;

use super::config::Config;
use super::sdk::SdkInfo;
use super::s3::S3;
use super::memdb::MemDb;
use super::utils::{ProgressIndicator, copy_with_progress, HumanDuration};
use super::{Result, ResultExt, ErrorKind};

/// Helper for synching
pub struct SyncOptions {
    pub user_facing: bool,
}

/// The main memdb stash type
pub struct MemDbStash {
    path: PathBuf,
    s3: S3,
    local_state: RwLock<Option<Arc<SdkSyncState>>>,
    memdbs: RwLock<HashMap<SdkInfo, Arc<MemDb<'static>>>>,
}

/// Information about a remotely available SDK
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct RemoteSdk {
    filename: String,
    info: SdkInfo,
    size: u64,
    etag: String,
}

#[derive(Serialize, Deserialize, Default, Debug, Clone)]
struct SdkSyncState {
    sdks: HashMap<String, RemoteSdk>,
    revision: Option<u64>,
}

/// Information about the health of the stash sync
#[derive(Debug)]
pub struct SyncStatus {
    remote_total: u32,
    missing: u32,
    different: u32,
    revision: u64,
}

impl RemoteSdk {
    /// Creates a remote SDK object from some information
    pub fn new(filename: String, info: SdkInfo, etag: String, size: u64) -> RemoteSdk {
        RemoteSdk {
            filename: filename,
            info: info,
            etag: etag,
            size: size,
        }
    }

    /// The remotely visible filename for the SDK
    pub fn filename(&self) -> &str {
        &self.filename
    }

    /// The local filename the SDK has in the stash folder
    pub fn local_filename(&self) -> &str {
        self.filename.trim_right_matches('z')
    }

    /// The size of the SDK in bytes
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Returns the SDK info
    pub fn info(&self) -> &SdkInfo {
        &self.info
    }
}

/// Iterator over the SDKs
pub type RemoteSdkIter<'a> = HashMapValuesIter<'a, String, RemoteSdk>;

impl SdkSyncState {
    pub fn get_sdk(&self, filename: &str) -> Option<&RemoteSdk> {
        self.sdks.get(filename)
    }

    pub fn update_sdk(&mut self, sdk: &RemoteSdk) {
        self.sdks.insert(sdk.local_filename().to_string(), sdk.clone());
    }

    pub fn sdks<'a>(&'a self) -> RemoteSdkIter<'a> {
        self.sdks.values()
    }
}

impl SyncStatus {

    /// Returns the lag (number of SDKs behind upstream)
    pub fn lag(&self) -> u32 {
        self.missing + self.different
    }

    /// Returns true if the local sync is still considered healthy
    pub fn is_healthy(&self) -> bool {
        let total = self.remote_total as f32;
        let lag = self.lag() as f32;
        lag / total < 0.10
    }

    /// Returns the revision of the stash
    pub fn revision(&self) -> u64 {
        self.revision
    }
}

impl Default for SyncOptions {
    fn default() -> SyncOptions {
        SyncOptions {
            user_facing: false,
        }
    }
}

impl MemDbStash {
    /// Opens a stash for a given config.
    pub fn new(config: &Config) -> Result<MemDbStash> {
        Ok(MemDbStash {
            path: config.get_symbol_dir()?.to_path_buf(),
            s3: S3::from_config(config)?,
            local_state: RwLock::new(None),
            memdbs: RwLock::new(HashMap::new()),
        })
    }

    fn get_local_sync_state_filename(&self) -> PathBuf {
        self.path.join("sync.state")
    }

    fn save_state(&self, new_state: &SdkSyncState, filename: &Path) -> Result<()> {
        let mut tmp_filename = filename.to_path_buf();
        tmp_filename.set_extension("tempstate");
        {
            let mut f = fs::File::create(&tmp_filename)?;
            serde_json::to_writer(&mut f, new_state)
                .chain_err(|| "Could not update sync state")?;
        }
        fs::rename(&tmp_filename, &filename)?;
        Ok(())
    }

    fn read_local_state(&self) -> Result<SdkSyncState> {
        let rv : SdkSyncState = match fs::File::open(&self.get_local_sync_state_filename()) {
            Ok(f) => serde_json::from_reader(f)
                .chain_err(|| "Parsing error on loading sync state")?,
            Err(err) => {
                if err.kind() == io::ErrorKind::NotFound {
                    Default::default()
                } else {
                    return Err(err).chain_err(|| "Error loading sync state");
                }
            }
        };
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(rv.clone()));
        Ok(rv)
    }

    fn get_local_state(&self) -> Result<Arc<SdkSyncState>> {
        if let Some(ref arc) = *self.local_state.read().unwrap() {
            return Ok(arc.clone());
        }
        self.read_local_state()?;
        Ok(self.local_state.read().unwrap().as_ref().unwrap().clone())
    }

    fn save_local_state(&self, new_state: &SdkSyncState) -> Result<()> {
        self.save_state(new_state, &self.get_local_sync_state_filename())?;
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(new_state.clone()));
        Ok(())
    }

    fn fetch_remote_state(&self) -> Result<SdkSyncState> {
        let mut sdks = HashMap::new();
        for remote_sdk in self.s3.list_upstream_sdks()? {
            sdks.insert(remote_sdk.local_filename().into(), remote_sdk);
        }
        Ok(SdkSyncState { sdks: sdks, revision: None })
    }

    fn update_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        // XXX: the progress bar here can stall out because we currently
        // need to buffer the download into memory in the s3 code :(
        let progress = if options.user_facing {
            ProgressIndicator::new(sdk.size() as usize)
        } else {
            info!("updating {}", sdk.info());
            ProgressIndicator::disabled()
        };
        let started = UTC::now();
        progress.set_message(&format!("Synchronizing {}", sdk.info()));
        let mut src = self.s3.download_sdk(sdk)?;
        let dst = fs::File::create(self.path.join(sdk.local_filename()))?;
        let mut dst = XzDecoder::new(dst);
        copy_with_progress(&progress, &mut src, &mut dst)?;
        progress.finish(&format!("Synchronized {}", sdk.info()));

        let duration = UTC::now() - started;
        if !options.user_facing {
            info!("updated {} in {}", sdk.info(), HumanDuration(duration));
        }
        Ok(())
    }

    fn remove_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        if options.user_facing {
            info!("  Deleting {}", sdk.info());
        } else {
            info!("removing {}", sdk.info());
        }
        if let Err(err) = fs::remove_file(self.path.join(sdk.local_filename())) {
            if err.kind() != io::ErrorKind::NotFound {
                return Err(err.into());
            }
        }
        Ok(())
    }

    /// Returns the current revision
    pub fn get_revision(&self) -> Result<u64> {
        Ok(self.read_local_state()?.revision.unwrap_or(0))
    }

    /// Checks the local stash against the server
    pub fn get_sync_status(&self) -> Result<SyncStatus> {
        let local_state = self.read_local_state()?;
        let remote_state = self.fetch_remote_state()?;

        let mut remote_total = 0;
        let mut missing = 0;
        let mut different = 0;

        for sdk in remote_state.sdks() {
            if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                if local_sdk != sdk {
                    different += 1;
                }
            } else {
                missing += 1;
            }
            remote_total += 1;
        }

        Ok(SyncStatus {
            remote_total: remote_total as u32,
            missing: missing as u32,
            different: different as u32,
            revision: local_state.revision.unwrap_or(0),
        })
    }

    /// Synchronize the local stash with the server
    pub fn sync(&self, options: SyncOptions) -> Result<()> {
        let mut local_state = self.read_local_state()?;
        let mut remote_state = self.fetch_remote_state()?;
        let started = UTC::now();
        let mut changed = false;

        let mut to_delete : HashSet<_> = HashSet::from_iter(
            local_state.sdks().map(|x| x.local_filename().to_string()));

        for sdk in remote_state.sdks() {
            if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                if local_sdk != sdk {
                    self.update_sdk(&sdk, &options)?;
                    changed = true;
                } else if options.user_facing {
                    println!("  ⸰ Unchanged {}", sdk.info());
                } else {
                    debug!("unchanged sdk {}", sdk.info());
                }
            } else {
                self.update_sdk(&sdk, &options)?;
                changed = true;
            }
            to_delete.remove(sdk.local_filename());
            local_state.update_sdk(&sdk);
            self.save_local_state(&local_state)?;
        }

        for local_filename in to_delete.iter() {
            if let Some(sdk) = local_state.get_sdk(local_filename) {
                self.remove_sdk(sdk, &options)?;
            }
        }

        let duration = UTC::now() - started;
        if options.user_facing {
            println!("Sync done in {}", HumanDuration(duration));
        } else if changed {
            info!("finished sync in {}", HumanDuration(duration));
        }

        // if we get this far, the remote state is indeed the local state.
        remote_state.revision = Some(local_state.revision.unwrap_or(0) + 1);
        self.save_local_state(&remote_state)?;

        Ok(())
    }

    /// Looks up an memdb by an SDK info if it's available.
    ///
    /// This returns a memdb wrapped in an arc as internally the system
    /// might try to unload the memdb if no longer needed.  If the MemDb
    /// does not exist, a `UnknownSdk` error is returned.
    pub fn get_memdb(&self, info: &SdkInfo) -> Result<Arc<MemDb<'static>>> {
        let local_state = self.get_local_state()?;

        // make sure we check in the local state first if the SDK exists.
        // if we go directly to the memdbs array or look at the file system
        // we might start to consider things that are not available yet or
        // not available any longer.
        if local_state.get_sdk(&info.memdb_filename()).is_some() {
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
            let memdb = MemDb::from_path(self.path.join(info.memdb_filename()))?;
            self.memdbs.write().unwrap().insert(info.clone(), Arc::new(memdb));
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
        }

        Err(ErrorKind::UnknownSdk.into())
    }

    /// Looks up an memdb by an SDK info as string if available.
    pub fn get_memdb_from_sdk_id(&self, sdk_id: &str) -> Result<Arc<MemDb<'static>>> {
        if let Some(sdk_info) = SdkInfo::from_filename(sdk_id) {
            self.get_memdb(&sdk_info)
        } else {
            Err(ErrorKind::UnknownSdk.into())
        }
    }
}
<commit_msg>Bump revision after every sync locally<commit_after>//! Provide access to locally cached memdb SDKs
//!
//! The `MemDbStash` pulls in remote SDKs from an S3 bucket and provides
//! access to it.  This is used by the symbol server to manage the local
//! cache and also to refer to memdb files that are mmap'ed in.
use std::fs;
use std::io;
use std::iter::FromIterator;
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Values as HashMapValuesIter;
use std::sync::{Arc, RwLock};

use serde_json;
use xz2::write::XzDecoder;
use chrono::UTC;

use super::config::Config;
use super::sdk::SdkInfo;
use super::s3::S3;
use super::memdb::MemDb;
use super::utils::{ProgressIndicator, copy_with_progress, HumanDuration};
use super::{Result, ResultExt, ErrorKind};

/// Helper for synching
pub struct SyncOptions {
    pub user_facing: bool,
}

/// The main memdb stash type
pub struct MemDbStash {
    path: PathBuf,
    s3: S3,
    local_state: RwLock<Option<Arc<SdkSyncState>>>,
    memdbs: RwLock<HashMap<SdkInfo, Arc<MemDb<'static>>>>,
}

/// Information about a remotely available SDK
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct RemoteSdk {
    filename: String,
    info: SdkInfo,
    size: u64,
    etag: String,
}

#[derive(Serialize, Deserialize, Default, Debug, Clone)]
struct SdkSyncState {
    sdks: HashMap<String, RemoteSdk>,
    revision: Option<u64>,
}

/// Information about the health of the stash sync
#[derive(Debug)]
pub struct SyncStatus {
    remote_total: u32,
    missing: u32,
    different: u32,
    revision: u64,
}

impl RemoteSdk {
    /// Creates a remote SDK object from some information
    pub fn new(filename: String, info: SdkInfo, etag: String, size: u64) -> RemoteSdk {
        RemoteSdk {
            filename: filename,
            info: info,
            etag: etag,
            size: size,
        }
    }

    /// The remotely visible filename for the SDK
    pub fn filename(&self) -> &str {
        &self.filename
    }

    /// The local filename the SDK has in the stash folder
    pub fn local_filename(&self) -> &str {
        self.filename.trim_right_matches('z')
    }

    /// The size of the SDK in bytes
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Returns the SDK info
    pub fn info(&self) -> &SdkInfo {
        &self.info
    }
}

/// Iterator over the SDKs
pub type RemoteSdkIter<'a> = HashMapValuesIter<'a, String, RemoteSdk>;

impl SdkSyncState {
    pub fn get_sdk(&self, filename: &str) -> Option<&RemoteSdk> {
        self.sdks.get(filename)
    }

    pub fn update_sdk(&mut self, sdk: &RemoteSdk) {
        self.sdks.insert(sdk.local_filename().to_string(), sdk.clone());
    }

    pub fn sdks<'a>(&'a self) -> RemoteSdkIter<'a> {
        self.sdks.values()
    }
}

impl SyncStatus {

    /// Returns the lag (number of SDKs behind upstream)
    pub fn lag(&self) -> u32 {
        self.missing + self.different
    }

    /// Returns true if the local sync is still considered healthy
    pub fn is_healthy(&self) -> bool {
        let total = self.remote_total as f32;
        let lag = self.lag() as f32;
        lag / total < 0.10
    }

    /// Returns the revision of the stash
    pub fn revision(&self) -> u64 {
        self.revision
    }
}

impl Default for SyncOptions {
    fn default() -> SyncOptions {
        SyncOptions {
            user_facing: false,
        }
    }
}

impl MemDbStash {
    /// Opens a stash for a given config.
    pub fn new(config: &Config) -> Result<MemDbStash> {
        Ok(MemDbStash {
            path: config.get_symbol_dir()?.to_path_buf(),
            s3: S3::from_config(config)?,
            local_state: RwLock::new(None),
            memdbs: RwLock::new(HashMap::new()),
        })
    }

    fn get_local_sync_state_filename(&self) -> PathBuf {
        self.path.join("sync.state")
    }

    fn save_state(&self, new_state: &SdkSyncState, filename: &Path) -> Result<()> {
        let mut tmp_filename = filename.to_path_buf();
        tmp_filename.set_extension("tempstate");
        {
            let mut f = fs::File::create(&tmp_filename)?;
            serde_json::to_writer(&mut f, new_state)
                .chain_err(|| "Could not update sync state")?;
        }
        fs::rename(&tmp_filename, &filename)?;
        Ok(())
    }

    fn read_local_state(&self) -> Result<SdkSyncState> {
        let rv : SdkSyncState = match fs::File::open(&self.get_local_sync_state_filename()) {
            Ok(f) => serde_json::from_reader(f)
                .chain_err(|| "Parsing error on loading sync state")?,
            Err(err) => {
                if err.kind() == io::ErrorKind::NotFound {
                    Default::default()
                } else {
                    return Err(err).chain_err(|| "Error loading sync state");
                }
            }
        };
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(rv.clone()));
        Ok(rv)
    }

    fn get_local_state(&self) -> Result<Arc<SdkSyncState>> {
        if let Some(ref arc) = *self.local_state.read().unwrap() {
            return Ok(arc.clone());
        }
        self.read_local_state()?;
        Ok(self.local_state.read().unwrap().as_ref().unwrap().clone())
    }

    fn save_local_state(&self, new_state: &SdkSyncState) -> Result<()> {
        self.save_state(new_state, &self.get_local_sync_state_filename())?;
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(new_state.clone()));
        Ok(())
    }

    fn fetch_remote_state(&self) -> Result<SdkSyncState> {
        let mut sdks = HashMap::new();
        for remote_sdk in self.s3.list_upstream_sdks()? {
            sdks.insert(remote_sdk.local_filename().into(), remote_sdk);
        }
        Ok(SdkSyncState { sdks: sdks, revision: None })
    }

    fn update_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        // XXX: the progress bar here can stall out because we currently
        // need to buffer the download into memory in the s3 code :(
        let progress = if options.user_facing {
            ProgressIndicator::new(sdk.size() as usize)
        } else {
            info!("updating {}", sdk.info());
            ProgressIndicator::disabled()
        };
        let started = UTC::now();
        progress.set_message(&format!("Synchronizing {}", sdk.info()));
        let mut src = self.s3.download_sdk(sdk)?;
        let dst = fs::File::create(self.path.join(sdk.local_filename()))?;
        let mut dst = XzDecoder::new(dst);
        copy_with_progress(&progress, &mut src, &mut dst)?;
        progress.finish(&format!("Synchronized {}", sdk.info()));

        let duration = UTC::now() - started;
        if !options.user_facing {
            info!("updated {} in {}", sdk.info(), HumanDuration(duration));
        }
        Ok(())
    }

    fn remove_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        if options.user_facing {
            info!("  Deleting {}", sdk.info());
        } else {
            info!("removing {}", sdk.info());
        }
        if let Err(err) = fs::remove_file(self.path.join(sdk.local_filename())) {
            if err.kind() != io::ErrorKind::NotFound {
                return Err(err.into());
            }
        }
        Ok(())
    }

    /// Returns the current revision
    pub fn get_revision(&self) -> Result<u64> {
        Ok(self.read_local_state()?.revision.unwrap_or(0))
    }

    /// Checks the local stash against the server
    pub fn get_sync_status(&self) -> Result<SyncStatus> {
        let local_state = self.read_local_state()?;
        let remote_state = self.fetch_remote_state()?;

        let mut remote_total = 0;
        let mut missing = 0;
        let mut different = 0;

        for sdk in remote_state.sdks() {
            if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                if local_sdk != sdk {
                    different += 1;
                }
            } else {
                missing += 1;
            }
            remote_total += 1;
        }

        Ok(SyncStatus {
            remote_total: remote_total as u32,
            missing: missing as u32,
            different: different as u32,
            revision: local_state.revision.unwrap_or(0),
        })
    }

    /// Synchronize the local stash with the server
    pub fn sync(&self, options: SyncOptions) -> Result<()> {
        let mut local_state = self.read_local_state()?;
        let mut remote_state = self.fetch_remote_state()?;
        let started = UTC::now();
        let mut changed = false;

        let mut to_delete : HashSet<_> = HashSet::from_iter(
            local_state.sdks().map(|x| x.local_filename().to_string()));

        for sdk in remote_state.sdks() {
            if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                if local_sdk != sdk {
                    self.update_sdk(&sdk, &options)?;
                    changed = true;
                } else if options.user_facing {
                    println!("  ⸰ Unchanged {}", sdk.info());
                } else {
                    debug!("unchanged sdk {}", sdk.info());
                }
            } else {
                self.update_sdk(&sdk, &options)?;
                changed = true;
            }
            to_delete.remove(sdk.local_filename());
            local_state.update_sdk(&sdk);
            local_state.revision = Some(local_state.revision.unwrap_or(0) + 1);
            self.save_local_state(&local_state)?;
        }

        for local_filename in to_delete.iter() {
            if let Some(sdk) = local_state.get_sdk(local_filename) {
                self.remove_sdk(sdk, &options)?;
            }
        }

        let duration = UTC::now() - started;
        if options.user_facing {
            println!("Sync done in {}", HumanDuration(duration));
        } else if changed {
            info!("finished sync in {}", HumanDuration(duration));
        }

        // if we get this far, the remote state is indeed the local state.
        remote_state.revision = Some(local_state.revision.unwrap_or(0) + 1);
        self.save_local_state(&remote_state)?;

        Ok(())
    }

    /// Looks up an memdb by an SDK info if it's available.
    ///
    /// This returns a memdb wrapped in an arc as internally the system
    /// might try to unload the memdb if no longer needed.  If the MemDb
    /// does not exist, a `UnknownSdk` error is returned.
    pub fn get_memdb(&self, info: &SdkInfo) -> Result<Arc<MemDb<'static>>> {
        let local_state = self.get_local_state()?;

        // make sure we check in the local state first if the SDK exists.
        // if we go directly to the memdbs array or look at the file system
        // we might start to consider things that are not available yet or
        // not available any longer.
        if local_state.get_sdk(&info.memdb_filename()).is_some() {
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
            let memdb = MemDb::from_path(self.path.join(info.memdb_filename()))?;
            self.memdbs.write().unwrap().insert(info.clone(), Arc::new(memdb));
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
        }

        Err(ErrorKind::UnknownSdk.into())
    }

    /// Looks up an memdb by an SDK info as string if available.
    pub fn get_memdb_from_sdk_id(&self, sdk_id: &str) -> Result<Arc<MemDb<'static>>> {
        if let Some(sdk_info) = SdkInfo::from_filename(sdk_id) {
            self.get_memdb(&sdk_info)
        } else {
            Err(ErrorKind::UnknownSdk.into())
        }
    }
}
<|endoftext|>


<commit_before>use std::fs;
use std::io::prelude::*;
use std::path;
use std::process;
use std::sync::mpsc::channel;
use std::thread;

use clap;
use cobalt::cobalt_model::files;
use hyper;
use hyper::server::{Server, Request, Response};
use hyper::uri::RequestUri;
use notify::{Watcher, RecursiveMode, raw_watcher};

use args;
use build;
use error::*;

pub fn watch_command_args() -> clap::App<'static, 'static> {
    clap::SubCommand::with_name("watch")
        .about("build, serve, and watch the project at the source dir")
        .args(&args::get_config_args())
        .arg(clap::Arg::with_name("port")
                 .short("P")
                 .long("port")
                 .value_name("INT")
                 .help("Port to serve from")
                 .default_value("3000")
                 .takes_value(true))
}

pub fn watch_command(matches: &clap::ArgMatches) -> Result<()> {
    let config = args::get_config(matches)?;
    let config = config.build()?;

    build::build(&config)?;

    // canonicalize is to ensure there is no question that `watcher`s paths come back safe for
    // Files::includes_file
    let source = path::Path::new(&config.source)
        .canonicalize()
        .chain_err(|| "Failed in processing source")?;
    let dest = path::Path::new(&config.destination).to_owned();

    // Be as broad as possible in what can cause a rebuild to
    // ensure we don't miss anything (normal file walks will miss
    // `_layouts`, etc).
    let mut site_files = files::FilesBuilder::new(&source)?;
    site_files.ignore_hidden(false)?;
    for line in &config.ignore {
        site_files.add_ignore(line.as_str())?;
    }
    let site_files = site_files.build()?;

    let port = matches.value_of("port").unwrap().to_string();
    thread::spawn(move || if serve(&dest, &port).is_err() {
                      process::exit(1)
                  });

    let (tx, rx) = channel();
    let mut watcher = raw_watcher(tx).chain_err(|| "Notify error")?;
    watcher
        .watch(&source, RecursiveMode::Recursive)
        .chain_err(|| "Notify error")?;
    info!("Watching {:?} for changes", &config.source);

    loop {
        let event = rx.recv().chain_err(|| "Notify error")?;
        let rebuild = if let Some(ref event_path) = event.path {
            if site_files.includes_file(event_path) {
                debug!("Page changed {:?}", event);
                true
            } else {
                trace!("Ignored file changed {:?}", event);
                false
            }
        } else {
            trace!("Assuming change {:?} is relevant", event);
            true
        };
        if rebuild {
            build::build(&config)?;
        }
    }
}

pub fn serve_command_args() -> clap::App<'static, 'static> {
    clap::SubCommand::with_name("serve")
        .about("build and serve the cobalt project at the source dir")
        .args(&args::get_config_args())
        .arg(clap::Arg::with_name("port")
                 .short("P")
                 .long("port")
                 .value_name("INT")
                 .help("Port to serve from")
                 .default_value("3000")
                 .takes_value(true))
}

pub fn serve_command(matches: &clap::ArgMatches) -> Result<()> {
    let config = args::get_config(matches)?;
    let config = config.build()?;

    build::build(&config)?;
    let port = matches.value_of("port").unwrap().to_string();
    let dest = path::Path::new(&config.destination);
    serve(dest, &port)?;

    Ok(())
}

fn static_file_handler(dest: &path::Path, req: Request, mut res: Response) -> Result<()> {
    // grab the requested path
    let mut req_path = match req.uri {
        RequestUri::AbsolutePath(p) => p,
        _ => {
            // return a 400 and exit from this request
            *res.status_mut() = hyper::status::StatusCode::BadRequest;
            let body = b"<h1> <center> 400: Bad request </center> </h1>";
            res.send(body)?;
            return Ok(());
        }
    };

    // strip off any querystrings so path.is_file() matches
    // and doesn't stick index.html on the end of the path
    // (querystrings often used for cachebusting)
    if let Some(position) = req_path.rfind('?') {
        req_path.truncate(position);
    }

    // find the path of the file in the local system
    // (this gets rid of the '/' in `p`, so the `join()` will not replace the
    // path)
    let path = dest.to_path_buf().join(&req_path[1..]);

    let serve_path = if path.is_file() {
        // try to point the serve path to `path` if it corresponds to a file
        path
    } else {
        // try to point the serve path into a "index.html" file in the requested
        // path
        path.join("index.html")
    };

    // if the request points to a file and it exists, read and serve it
    if serve_path.exists() {
        let mut file = fs::File::open(serve_path)?;

        // buffer to store the file
        let mut buffer: Vec<u8> = vec![];

        file.read_to_end(&mut buffer)?;

        res.send(&buffer)?;
    } else {
        // return a 404 status
        *res.status_mut() = hyper::status::StatusCode::NotFound;

        // write a simple body for the 404 page
        let body = b"<h1> <center> 404: Page not found </center> </h1>";

        res.send(body)?;
    }

    Ok(())
}

fn serve(dest: &path::Path, port: &str) -> Result<()> {
    info!("Serving {:?} through static file server", dest);

    let ip = format!("127.0.0.1:{}", port);
    info!("Server Listening on {}", &ip);
    info!("Ctrl-c to stop the server");

    // attempts to create a server
    let http_server = match Server::http(&*ip) {
        Ok(server) => server,
        Err(e) => {
            error!("{}", e);
            return Err(e.into());
        }
    };

    // need a clone because of closure's lifetime
    let dest_clone = dest.to_owned();

    // bind the handle function and start serving
    if let Err(e) = http_server.handle(move |req: Request, res: Response| if let Err(e) =
        static_file_handler(&dest_clone, req, res) {
        error!("{}", e);
        process::exit(1);
    }) {
        error!("{}", e);
        return Err(e.into());
    };

    Ok(())
}
<commit_msg>fix(watch): Don't stop on error<commit_after>use std::fs;
use std::io::prelude::*;
use std::path;
use std::process;
use std::sync::mpsc::channel;
use std::thread;

use clap;
use cobalt::cobalt_model::files;
use error_chain::ChainedError;
use hyper;
use hyper::server::{Server, Request, Response};
use hyper::uri::RequestUri;
use notify::{Watcher, RecursiveMode, raw_watcher};

use args;
use build;
use error::*;

pub fn watch_command_args() -> clap::App<'static, 'static> {
    clap::SubCommand::with_name("watch")
        .about("build, serve, and watch the project at the source dir")
        .args(&args::get_config_args())
        .arg(clap::Arg::with_name("port")
                 .short("P")
                 .long("port")
                 .value_name("INT")
                 .help("Port to serve from")
                 .default_value("3000")
                 .takes_value(true))
}

pub fn watch_command(matches: &clap::ArgMatches) -> Result<()> {
    let config = args::get_config(matches)?;
    let config = config.build()?;

    build::build(&config)?;

    // canonicalize is to ensure there is no question that `watcher`s paths come back safe for
    // Files::includes_file
    let source = path::Path::new(&config.source)
        .canonicalize()
        .chain_err(|| "Failed in processing source")?;
    let dest = path::Path::new(&config.destination).to_owned();

    // Be as broad as possible in what can cause a rebuild to
    // ensure we don't miss anything (normal file walks will miss
    // `_layouts`, etc).
    let mut site_files = files::FilesBuilder::new(&source)?;
    site_files.ignore_hidden(false)?;
    for line in &config.ignore {
        site_files.add_ignore(line.as_str())?;
    }
    let site_files = site_files.build()?;

    let port = matches.value_of("port").unwrap().to_string();
    thread::spawn(move || if serve(&dest, &port).is_err() {
                      process::exit(1)
                  });

    let (tx, rx) = channel();
    let mut watcher = raw_watcher(tx).chain_err(|| "Notify error")?;
    watcher
        .watch(&source, RecursiveMode::Recursive)
        .chain_err(|| "Notify error")?;
    info!("Watching {:?} for changes", &config.source);

    loop {
        let event = rx.recv().chain_err(|| "Notify error")?;
        let rebuild = if let Some(ref event_path) = event.path {
            if site_files.includes_file(event_path) {
                debug!("Page changed {:?}", event);
                true
            } else {
                trace!("Ignored file changed {:?}", event);
                false
            }
        } else {
            trace!("Assuming change {:?} is relevant", event);
            true
        };
        if rebuild {
            let result = build::build(&config);
            if let Err(fail) = result {
                error!("build failed\n{}", fail.display_chain());
            }
        }
    }
}

pub fn serve_command_args() -> clap::App<'static, 'static> {
    clap::SubCommand::with_name("serve")
        .about("build and serve the cobalt project at the source dir")
        .args(&args::get_config_args())
        .arg(clap::Arg::with_name("port")
                 .short("P")
                 .long("port")
                 .value_name("INT")
                 .help("Port to serve from")
                 .default_value("3000")
                 .takes_value(true))
}

pub fn serve_command(matches: &clap::ArgMatches) -> Result<()> {
    let config = args::get_config(matches)?;
    let config = config.build()?;

    build::build(&config)?;
    let port = matches.value_of("port").unwrap().to_string();
    let dest = path::Path::new(&config.destination);
    serve(dest, &port)?;

    Ok(())
}

fn static_file_handler(dest: &path::Path, req: Request, mut res: Response) -> Result<()> {
    // grab the requested path
    let mut req_path = match req.uri {
        RequestUri::AbsolutePath(p) => p,
        _ => {
            // return a 400 and exit from this request
            *res.status_mut() = hyper::status::StatusCode::BadRequest;
            let body = b"<h1> <center> 400: Bad request </center> </h1>";
            res.send(body)?;
            return Ok(());
        }
    };

    // strip off any querystrings so path.is_file() matches
    // and doesn't stick index.html on the end of the path
    // (querystrings often used for cachebusting)
    if let Some(position) = req_path.rfind('?') {
        req_path.truncate(position);
    }

    // find the path of the file in the local system
    // (this gets rid of the '/' in `p`, so the `join()` will not replace the
    // path)
    let path = dest.to_path_buf().join(&req_path[1..]);

    let serve_path = if path.is_file() {
        // try to point the serve path to `path` if it corresponds to a file
        path
    } else {
        // try to point the serve path into a "index.html" file in the requested
        // path
        path.join("index.html")
    };

    // if the request points to a file and it exists, read and serve it
    if serve_path.exists() {
        let mut file = fs::File::open(serve_path)?;

        // buffer to store the file
        let mut buffer: Vec<u8> = vec![];

        file.read_to_end(&mut buffer)?;

        res.send(&buffer)?;
    } else {
        // return a 404 status
        *res.status_mut() = hyper::status::StatusCode::NotFound;

        // write a simple body for the 404 page
        let body = b"<h1> <center> 404: Page not found </center> </h1>";

        res.send(body)?;
    }

    Ok(())
}

fn serve(dest: &path::Path, port: &str) -> Result<()> {
    info!("Serving {:?} through static file server", dest);

    let ip = format!("127.0.0.1:{}", port);
    info!("Server Listening on {}", &ip);
    info!("Ctrl-c to stop the server");

    // attempts to create a server
    let http_server = match Server::http(&*ip) {
        Ok(server) => server,
        Err(e) => {
            error!("{}", e);
            return Err(e.into());
        }
    };

    // need a clone because of closure's lifetime
    let dest_clone = dest.to_owned();

    // bind the handle function and start serving
    if let Err(e) = http_server.handle(move |req: Request, res: Response| if let Err(e) =
        static_file_handler(&dest_clone, req, res) {
        error!("{}", e);
        process::exit(1);
    }) {
        error!("{}", e);
        return Err(e.into());
    };

    Ok(())
}
<|endoftext|>


<commit_before><commit_msg>Add example using custom colors<commit_after><|endoftext|>