<commit_before>use clap::{App, Arg};
use regex::Regex;

fn integer_value(v: String) -> Result<(), String> {
  v.parse::<u16>().map(|_| ()).map_err(|e| format!("'{}' is not a valid port value: {}", v, e) )
}

pub(crate) fn setup_app<'a, 'b>(program: String, version: &'b str) -> App<'a, 'b> {
  App::new(program)
    .version(version)
    .about("Standalone Pact verifier")
    .version_short("v")
    .arg(Arg::with_name("loglevel")
      .short("l")
      .long("loglevel")
      .takes_value(true)
      .use_delimiter(false)
      .possible_values(&["error", "warn", "info", "debug", "trace", "none"])
      .help("Log level (defaults to warn)"))
    .arg(Arg::with_name("file")
      .short("f")
      .long("file")
      .required_unless_one(&["dir", "url", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("Pact file to verify (can be repeated)"))
    .arg(Arg::with_name("dir")
      .short("d")
      .long("dir")
      .required_unless_one(&["file", "url", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("Directory of pact files to verify (can be repeated)"))
    .arg(Arg::with_name("url")
      .short("u")
      .long("url")
      .required_unless_one(&["file", "dir", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of pact file to verify (can be repeated)"))
    .arg(Arg::with_name("broker-url")
      .short("b")
      .long("broker-url")
      .env("PACT_BROKER_BASE_URL")
      .required_unless_one(&["file", "dir", "url"])
      .requires("provider-name")
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of the pact broker to fetch pacts from to verify (requires the provider name parameter)"))
    .arg(Arg::with_name("hostname")
      .short("h")
      .long("hostname")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider hostname (defaults to localhost)"))
    .arg(Arg::with_name("port")
      .short("p")
      .long("port")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider port (defaults to protocol default 80/443)")
      .validator(integer_value))
    .arg(Arg::with_name("scheme")
      .long("scheme")
      .takes_value(true)
      .possible_values(&["http", "https"])
      .default_value("http")
      .help("Provider URI scheme (defaults to http)"))
    .arg(Arg::with_name("provider-name")
      .short("n")
      .long("provider-name")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider name (defaults to provider)"))
    .arg(Arg::with_name("state-change-url")
      .short("s")
      .long("state-change-url")
      .takes_value(true)
      .use_delimiter(false)
      .help("URL to post state change requests to"))
    .arg(Arg::with_name("state-change-as-query")
      .long("state-change-as-query")
      .help("State change request data will be sent as query parameters instead of in the request body"))
    .arg(Arg::with_name("state-change-teardown")
      .long("state-change-teardown")
      .help("State change teardown requests are to be made after each interaction"))
    .arg(Arg::with_name("filter-description")
      .long("filter-description")
      .env("PACT_DESCRIPTION")
      .takes_value(true)
      .use_delimiter(false)
      .validator(|val| Regex::new(&val)
        .map(|_| ())
        .map_err(|err| format!("'{}' is an invalid filter value: {}", val, err)))
      .help("Only validate interactions whose descriptions match this filter"))
    .arg(Arg::with_name("filter-state")
      .long("filter-state")
      .env("PACT_PROVIDER_STATE")
      .takes_value(true)
      .use_delimiter(false)
      .conflicts_with("filter-no-state")
      .validator(|val| Regex::new(&val)
        .map(|_| ())
        .map_err(|err| format!("'{}' is an invalid filter value: {}", val, err)))
      .help("Only validate interactions whose provider states match this filter"))
    .arg(Arg::with_name("filter-no-state")
      .long("filter-no-state")
      .env("PACT_PROVIDER_NO_STATE")
      .conflicts_with("filter-state")
      .help("Only validate interactions that have no defined provider state"))
    .arg(Arg::with_name("filter-consumer")
      .short("c")
      .long("filter-consumer")
      .takes_value(true)
      .multiple(true)
      .empty_values(false)
      .help("Consumer name to filter the pacts to be verified (can be repeated)"))
    .arg(Arg::with_name("user")
      .long("user")
      .env("PACT_BROKER_USERNAME")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("token")
      .help("Username to use when fetching pacts from URLS"))
      .arg(Arg::with_name("password")
      .long("password")
      .env("PACT_BROKER_PASSWORD")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("token")
      .help("Password to use when fetching pacts from URLS"))
      .arg(Arg::with_name("token")
      .short("t")
      .long("token")
      .env("PACT_BROKER_TOKEN")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("user")
      .help("Bearer token to use when fetching pacts from URLS"))
    .arg(Arg::with_name("publish")
      .long("publish")
      .requires("broker-url")
      .requires("provider-version")
      .help("Enables publishing of verification results back to the Pact Broker. Requires the broker-url and provider-version parameters."))
    .arg(Arg::with_name("provider-version")
      .long("provider-version")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .help("Provider version that is being verified. This is required when publishing results."))
    .arg(Arg::with_name("build-url")
      .long("build-url")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of the build to associate with the published verification results."))
    .arg(Arg::with_name("provider-tags")
      .long("provider-tags")
      .takes_value(true)
      .use_delimiter(true)
      .empty_values(false)
      .help("Provider tags to use when publishing results. Accepts comma-separated values."))
    .arg(Arg::with_name("base-path")
      .long("base-path")
      .takes_value(true)
      .use_delimiter(false)
      .empty_values(false)
      .help("Base path to add to all requests"))
    .arg(Arg::with_name("consumer-version-tags")
      .long("consumer-version-tags")
      .takes_value(true)
      .use_delimiter(true)
      .empty_values(false)
      .requires("broker-url")
      .conflicts_with("consumer-version-selectors")
      .help("Consumer tags to use when fetching pacts from the Broker. Accepts comma-separated values."))
    .arg(Arg::with_name("consumer-version-selectors")
      .long("consumer-version-selectors")
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .requires("broker-url")
      .conflicts_with("consumer-version-tags")
      .help("Consumer version selectors to use when fetching pacts from the Broker. Accepts a JSON string as per https://docs.pact.io/pact_broker/advanced_topics/consumer_version_selectors/"))
    .arg(Arg::with_name("disable-ssl-verification")
      .long("disable-ssl-verification")
      .takes_value(false)
      .help("Disables validation of SSL certificates"))
    .arg(Arg::with_name("enable-pending")
      .long("enable-pending")
      .requires("broker-url")
      .help("Enables Pending Pacts"))
    .arg(Arg::with_name("include-wip-pacts-since")
      .long("include-wip-pacts-since")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .requires("broker-url")
      .help("Allow pacts that don't match given consumer selectors (or tags) to  be verified, without causing the overall task to fail. For more information, see https://pact.io/wip"))
    .arg(Arg::with_name("request-timeout")
      .long("request-timeout")
      .takes_value(true)
      .empty_values(false)
      .help("Sets the HTTP request timeout in milliseconds for requests to the target API and for state change requests."))
    }

#[cfg(test)]
mod test {

  use quickcheck::{TestResult, quickcheck};
  use rand::Rng;
  use super::integer_value;
  use expectest::prelude::*;
  use expectest::expect;
  use pact_matching::s;

  #[test]
  fn validates_integer_value() {
    fn prop(s: String) -> TestResult {
      let mut rng = ::rand::thread_rng();
      if rng.gen() && s.chars().any(|ch| !ch.is_numeric()) {
        TestResult::discard()
      } else {
        let validation = integer_value(s.clone());
        match validation {
          Ok(_) => TestResult::from_bool(!s.is_empty() && s.chars().all(|ch| ch.is_numeric() )),
          Err(_) => TestResult::from_bool(s.is_empty() || s.chars().find(|ch| !ch.is_numeric() ).is_some())
        }
      }
    }
    quickcheck(prop as fn(_) -> _);

    expect!(integer_value(s!("1234"))).to(be_ok());
    expect!(integer_value(s!("1234x"))).to(be_err());
  }
}
<commit_msg>wip: add validator to request timeout<commit_after>use clap::{App, Arg};
use regex::Regex;

fn integer_value(v: String) -> Result<(), String> {
  v.parse::<u16>().map(|_| ()).map_err(|e| format!("'{}' is not a valid port value: {}", v, e) )
}

pub(crate) fn setup_app<'a, 'b>(program: String, version: &'b str) -> App<'a, 'b> {
  App::new(program)
    .version(version)
    .about("Standalone Pact verifier")
    .version_short("v")
    .arg(Arg::with_name("loglevel")
      .short("l")
      .long("loglevel")
      .takes_value(true)
      .use_delimiter(false)
      .possible_values(&["error", "warn", "info", "debug", "trace", "none"])
      .help("Log level (defaults to warn)"))
    .arg(Arg::with_name("file")
      .short("f")
      .long("file")
      .required_unless_one(&["dir", "url", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("Pact file to verify (can be repeated)"))
    .arg(Arg::with_name("dir")
      .short("d")
      .long("dir")
      .required_unless_one(&["file", "url", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("Directory of pact files to verify (can be repeated)"))
    .arg(Arg::with_name("url")
      .short("u")
      .long("url")
      .required_unless_one(&["file", "dir", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of pact file to verify (can be repeated)"))
    .arg(Arg::with_name("broker-url")
      .short("b")
      .long("broker-url")
      .env("PACT_BROKER_BASE_URL")
      .required_unless_one(&["file", "dir", "url"])
      .requires("provider-name")
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of the pact broker to fetch pacts from to verify (requires the provider name parameter)"))
    .arg(Arg::with_name("hostname")
      .short("h")
      .long("hostname")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider hostname (defaults to localhost)"))
    .arg(Arg::with_name("port")
      .short("p")
      .long("port")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider port (defaults to protocol default 80/443)")
      .validator(integer_value))
    .arg(Arg::with_name("scheme")
      .long("scheme")
      .takes_value(true)
      .possible_values(&["http", "https"])
      .default_value("http")
      .help("Provider URI scheme (defaults to http)"))
    .arg(Arg::with_name("provider-name")
      .short("n")
      .long("provider-name")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider name (defaults to provider)"))
    .arg(Arg::with_name("state-change-url")
      .short("s")
      .long("state-change-url")
      .takes_value(true)
      .use_delimiter(false)
      .help("URL to post state change requests to"))
    .arg(Arg::with_name("state-change-as-query")
      .long("state-change-as-query")
      .help("State change request data will be sent as query parameters instead of in the request body"))
    .arg(Arg::with_name("state-change-teardown")
      .long("state-change-teardown")
      .help("State change teardown requests are to be made after each interaction"))
    .arg(Arg::with_name("filter-description")
      .long("filter-description")
      .env("PACT_DESCRIPTION")
      .takes_value(true)
      .use_delimiter(false)
      .validator(|val| Regex::new(&val)
        .map(|_| ())
        .map_err(|err| format!("'{}' is an invalid filter value: {}", val, err)))
      .help("Only validate interactions whose descriptions match this filter"))
    .arg(Arg::with_name("filter-state")
      .long("filter-state")
      .env("PACT_PROVIDER_STATE")
      .takes_value(true)
      .use_delimiter(false)
      .conflicts_with("filter-no-state")
      .validator(|val| Regex::new(&val)
        .map(|_| ())
        .map_err(|err| format!("'{}' is an invalid filter value: {}", val, err)))
      .help("Only validate interactions whose provider states match this filter"))
    .arg(Arg::with_name("filter-no-state")
      .long("filter-no-state")
      .env("PACT_PROVIDER_NO_STATE")
      .conflicts_with("filter-state")
      .help("Only validate interactions that have no defined provider state"))
    .arg(Arg::with_name("filter-consumer")
      .short("c")
      .long("filter-consumer")
      .takes_value(true)
      .multiple(true)
      .empty_values(false)
      .help("Consumer name to filter the pacts to be verified (can be repeated)"))
    .arg(Arg::with_name("user")
      .long("user")
      .env("PACT_BROKER_USERNAME")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("token")
      .help("Username to use when fetching pacts from URLS"))
      .arg(Arg::with_name("password")
      .long("password")
      .env("PACT_BROKER_PASSWORD")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("token")
      .help("Password to use when fetching pacts from URLS"))
      .arg(Arg::with_name("token")
      .short("t")
      .long("token")
      .env("PACT_BROKER_TOKEN")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("user")
      .help("Bearer token to use when fetching pacts from URLS"))
    .arg(Arg::with_name("publish")
      .long("publish")
      .requires("broker-url")
      .requires("provider-version")
      .help("Enables publishing of verification results back to the Pact Broker. Requires the broker-url and provider-version parameters."))
    .arg(Arg::with_name("provider-version")
      .long("provider-version")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .help("Provider version that is being verified. This is required when publishing results."))
    .arg(Arg::with_name("build-url")
      .long("build-url")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of the build to associate with the published verification results."))
    .arg(Arg::with_name("provider-tags")
      .long("provider-tags")
      .takes_value(true)
      .use_delimiter(true)
      .empty_values(false)
      .help("Provider tags to use when publishing results. Accepts comma-separated values."))
    .arg(Arg::with_name("base-path")
      .long("base-path")
      .takes_value(true)
      .use_delimiter(false)
      .empty_values(false)
      .help("Base path to add to all requests"))
    .arg(Arg::with_name("consumer-version-tags")
      .long("consumer-version-tags")
      .takes_value(true)
      .use_delimiter(true)
      .empty_values(false)
      .requires("broker-url")
      .conflicts_with("consumer-version-selectors")
      .help("Consumer tags to use when fetching pacts from the Broker. Accepts comma-separated values."))
    .arg(Arg::with_name("consumer-version-selectors")
      .long("consumer-version-selectors")
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .requires("broker-url")
      .conflicts_with("consumer-version-tags")
      .help("Consumer version selectors to use when fetching pacts from the Broker. Accepts a JSON string as per https://docs.pact.io/pact_broker/advanced_topics/consumer_version_selectors/"))
    .arg(Arg::with_name("disable-ssl-verification")
      .long("disable-ssl-verification")
      .takes_value(false)
      .help("Disables validation of SSL certificates"))
    .arg(Arg::with_name("enable-pending")
      .long("enable-pending")
      .requires("broker-url")
      .help("Enables Pending Pacts"))
    .arg(Arg::with_name("include-wip-pacts-since")
      .long("include-wip-pacts-since")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .requires("broker-url")
      .help("Allow pacts that don't match given consumer selectors (or tags) to  be verified, without causing the overall task to fail. For more information, see https://pact.io/wip"))
    .arg(Arg::with_name("request-timeout")
      .long("request-timeout")
      .takes_value(true)
      .empty_values(false)
      .validator(integer_value)
      .help("Sets the HTTP request timeout in milliseconds for requests to the target API and for state change requests."))
    }

#[cfg(test)]
mod test {

  use quickcheck::{TestResult, quickcheck};
  use rand::Rng;
  use super::integer_value;
  use expectest::prelude::*;
  use expectest::expect;
  use pact_matching::s;

  #[test]
  fn validates_integer_value() {
    fn prop(s: String) -> TestResult {
      let mut rng = ::rand::thread_rng();
      if rng.gen() && s.chars().any(|ch| !ch.is_numeric()) {
        TestResult::discard()
      } else {
        let validation = integer_value(s.clone());
        match validation {
          Ok(_) => TestResult::from_bool(!s.is_empty() && s.chars().all(|ch| ch.is_numeric() )),
          Err(_) => TestResult::from_bool(s.is_empty() || s.chars().find(|ch| !ch.is_numeric() ).is_some())
        }
      }
    }
    quickcheck(prop as fn(_) -> _);

    expect!(integer_value(s!("1234"))).to(be_ok());
    expect!(integer_value(s!("1234x"))).to(be_err());
  }
}
<|endoftext|>


<commit_before>#![feature(rustc_private)]

// NOTE: For the example to compile, you will need to first run the following:
//     rustup component add rustc-dev llvm-tools-preview

// version: 1.53.0-nightly (9b0edb7fd 2021-03-27)

extern crate rustc_ast_pretty;
extern crate rustc_error_codes;
extern crate rustc_errors;
extern crate rustc_hash;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_session;
extern crate rustc_span;

use rustc_ast_pretty::pprust::item_to_string;
use rustc_errors::registry;
use rustc_session::config;
use rustc_span::source_map;
use std::path;
use std::process;
use std::str;

fn main() {
    let out = process::Command::new("rustc")
        .arg("--print=sysroot")
        .current_dir(".")
        .output()
        .unwrap();
    let sysroot = str::from_utf8(&out.stdout).unwrap().trim();
    let config = rustc_interface::Config {
        opts: config::Options {
            maybe_sysroot: Some(path::PathBuf::from(sysroot)),
            ..config::Options::default()
        },
        input: config::Input::Str {
            name: source_map::FileName::Custom("main.rs".to_string()),
            input: "fn main() { let message = \"Hello, world!\"; println!(\"{}\", message); }"
                .to_string(),
        },
        diagnostic_output: rustc_session::DiagnosticOutput::Default,
        crate_cfg: rustc_hash::FxHashSet::default(),
        input_path: None,
        output_dir: None,
        output_file: None,
        file_loader: None,
        stderr: None,
        lint_caps: rustc_hash::FxHashMap::default(),
        parse_sess_created: None,
        register_lints: None,
        override_queries: None,
        make_codegen_backend: None,
        registry: registry::Registry::new(&rustc_error_codes::DIAGNOSTICS),
    };
    rustc_interface::run_compiler(config, |compiler| {
        compiler.enter(|queries| {
            // TODO: add this to -Z unpretty
            let ast_krate = queries.parse().unwrap().take();
            for item in ast_krate.items {
                println!("{}", item_to_string(&item));
            }

            // Analyze the crate and inspect the types under the cursor.
            queries.global_ctxt().unwrap().take().enter(|tcx| {
                // Every compilation contains a single crate.
                let hir_krate = tcx.hir().krate();
                // Iterate over the top-level items in the crate, looking for the main function.
                for (_, item) in &hir_krate.items {
                    // Use pattern-matching to find a specific node inside the main function.
                    if let rustc_hir::ItemKind::Fn(_, _, body_id) = item.kind {
                        let expr = &tcx.hir().body(body_id).value;
                        if let rustc_hir::ExprKind::Block(block, _) = expr.kind {
                            if let rustc_hir::StmtKind::Local(local) = block.stmts[0].kind {
                                if let Some(expr) = local.init {
                                    let hir_id = expr.hir_id; // hir_id identifies the string "Hello, world!"
                                    let def_id = tcx.hir().local_def_id(item.hir_id()); // def_id identifies the main function
                                    let ty = tcx.typeck(def_id).node_type(hir_id);
                                    println!("{:?}: {:?}", expr, ty); // prints expr(HirId { owner: DefIndex(3), local_id: 4 }: "Hello, world!"): &'static str
                                }
                            }
                        }
                    }
                }
            })
        });
    });
}
<commit_msg>make it compile with 1.56.0<commit_after>#![feature(rustc_private)]

// NOTE: For the example to compile, you will need to first run the following:
//     rustup component add rustc-dev llvm-tools-preview

// version: 1.53.0-nightly (9b0edb7fd 2021-03-27)

extern crate rustc_ast_pretty;
extern crate rustc_error_codes;
extern crate rustc_errors;
extern crate rustc_hash;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_session;
extern crate rustc_span;

use rustc_ast_pretty::pprust::item_to_string;
use rustc_errors::registry;
use rustc_session::config;
use rustc_span::source_map;
use std::path;
use std::process;
use std::str;

fn main() {
    let out = process::Command::new("rustc")
        .arg("--print=sysroot")
        .current_dir(".")
        .output()
        .unwrap();
    let sysroot = str::from_utf8(&out.stdout).unwrap().trim();
    let config = rustc_interface::Config {
        opts: config::Options {
            maybe_sysroot: Some(path::PathBuf::from(sysroot)),
            ..config::Options::default()
        },
        input: config::Input::Str {
            name: source_map::FileName::Custom("main.rs".to_string()),
            input: "fn main() { let message = \"Hello, world!\"; println!(\"{}\", message); }"
                .to_string(),
        },
        diagnostic_output: rustc_session::DiagnosticOutput::Default,
        crate_cfg: rustc_hash::FxHashSet::default(),
        input_path: None,
        output_dir: None,
        output_file: None,
        file_loader: None,
        stderr: None,
        lint_caps: rustc_hash::FxHashMap::default(),
        parse_sess_created: None,
        register_lints: None,
        override_queries: None,
        make_codegen_backend: None,
        registry: registry::Registry::new(&rustc_error_codes::DIAGNOSTICS),
    };
    rustc_interface::run_compiler(config, |compiler| {
        compiler.enter(|queries| {
            // TODO: add this to -Z unpretty
            let ast_krate = queries.parse().unwrap().take();
            for item in ast_krate.items {
                println!("{}", item_to_string(&item));
            }

            // Analyze the crate and inspect the types under the cursor.
            queries.global_ctxt().unwrap().take().enter(|tcx| {
                // Every compilation contains a single crate.
                let hir_krate = tcx.hir();
                // Iterate over the top-level items in the crate, looking for the main function.
                for item in hir_krate.items() {
                    // Use pattern-matching to find a specific node inside the main function.
                    if let rustc_hir::ItemKind::Fn(_, _, body_id) = item.kind {
                        let expr = &tcx.hir().body(body_id).value;
                        if let rustc_hir::ExprKind::Block(block, _) = expr.kind {
                            if let rustc_hir::StmtKind::Local(local) = block.stmts[0].kind {
                                if let Some(expr) = local.init {
                                    let hir_id = expr.hir_id; // hir_id identifies the string "Hello, world!"
                                    let def_id = tcx.hir().local_def_id(item.hir_id()); // def_id identifies the main function
                                    let ty = tcx.typeck(def_id).node_type(hir_id);
                                    println!("{:?}: {:?}", expr, ty); // prints expr(HirId { owner: DefIndex(3), local_id: 4 }: "Hello, world!"): &'static str
                                }
                            }
                        }
                    }
                }
            })
        });
    });
}
<|endoftext|>


<commit_before>use iron::typemap::Key;
use uuid::Uuid;
use serde_json;
use std::iter;
use r2d2;
use r2d2_postgres;
use postgres;
use std::convert::From;
use std::collections::HashMap;
use std::fmt;

error_chain! {
    errors {
        UniqueViolation
        NoTestQuery(cmd: String)
        NoTestExecute(cmd: String)
    }
    foreign_links {
        R2D2Error(r2d2::Error);
        PostgresError(postgres::Error);
    }
}

pub struct HexSlice<'a>(&'a [u8]);

impl<'a> HexSlice<'a> {
    pub fn new<T>(data: &'a T) -> HexSlice<'a>
        where T: ?Sized + AsRef<[u8]> + 'a
    {
        HexSlice(data.as_ref())
    }
}

impl<'a> fmt::Display for HexSlice<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for byte in self.0 {
            write!(f, "{:02X}", byte)?;
        }
        Ok(())
    }
}

pub trait FromSql {}
impl FromSql for u32 {}
impl FromSql for Uuid {}
impl FromSql for String {}
impl <'a> FromSql for &'a str {}
impl FromSql for Vec<u8> {}
impl FromSql for serde_json::Value {}

#[derive(Debug, Clone)]
pub enum SqlValue {
    U32(u32),
    UUID(Uuid),
    String(String),
    U8Bytes(Vec<u8>),
    JSON(serde_json::Value),
}

impl From<u32> for SqlValue {
    fn from(i: u32) -> SqlValue { SqlValue::U32(i) }
}

#[derive(Hash, Eq, PartialEq, Clone, Debug)]
pub enum ValueIndex
{
    U32(u32),
    String(String),
}

impl<'a> From<&'a str> for ValueIndex {
    fn from(s: &str) -> ValueIndex { ValueIndex::String(String::from(s))}
}

pub trait RowIndex {
    fn val(&self) -> ValueIndex;
}
impl RowIndex for u32 {
    fn val(&self) -> ValueIndex { ValueIndex::U32(*self)}
}
impl RowIndex for String {
    fn val(&self) -> ValueIndex { ValueIndex::String(self.clone())}
}
impl <'a> RowIndex for &'a str {
    fn val(&self) -> ValueIndex { ValueIndex::String(String::from(*self))}
}

#[derive(Debug, Clone)]
pub struct TestRow
{
    data: HashMap<ValueIndex, SqlValue>
}

pub trait GetRow<T> {
    fn get<R>(&self, id: R) -> T where T: FromSql, R: RowIndex + fmt::Display;
}

impl TestRow {
    pub fn new() -> TestRow {
        TestRow{data: HashMap::new()}
    }

    pub fn insert<K, V>(&mut self, k: K, v: V) -> Option<SqlValue>
        where K: Into<ValueIndex>, V: Into<SqlValue> {
        self.data.insert(k.into(), v.into())
    }
}

macro_rules! get_row {
    ($type: ty, $kind: path) => (impl<'a> GetRow<$type> for TestRow {
    fn get<R>(&self, id: R) -> $type where R: RowIndex + fmt::Display {
        if !self.data.contains_key(&id.val()) {
            panic!(format!("Can't find key {} in row", id));
        }
        match self.data[&id.val()] {
            $kind(ref val) => val.clone(),
            _ => panic!()
        }
    }
})}

get_row!(u32, SqlValue::U32);
get_row!(Uuid, SqlValue::UUID);
get_row!(String, SqlValue::String);
get_row!(serde_json::Value, SqlValue::JSON);
get_row!(Vec<u8>, SqlValue::U8Bytes);

pub enum Row<'a> {
    Postgres(postgres::rows::Row<'a>),
    Test(&'a TestRow)
}
impl<'a> Row<'a> {
    pub fn get<T, R>(&self, id: R) -> T
    where
        T: FromSql + postgres::types::FromSql,
        R: RowIndex + postgres::rows::RowIndex + fmt::Display + fmt::Debug,
        TestRow: GetRow<T> {
        match self {
            &Row::Postgres(ref rows) => {
                rows.get(id)
            }
            &Row::Test(ref rows) => {
                rows.get(id)
            }
        }
    }
    pub fn get_opt<T, R>(&self, id: R) -> Option<Result<T>> where T: FromSql, R: RowIndex {
        unimplemented!();
    }
}

pub struct TestRowIterator<'a> {
    rows: &'a Vec<TestRow>,
    location: usize
}

impl<'a> TestRowIterator<'a> {
    fn new(r: &'a Vec<TestRow>) -> TestRowIterator<'a> {
        TestRowIterator {
            rows: r,
            location: 0
        }
    }

    fn next(&mut self) -> Option<&'a TestRow> {
        self.location +=1;
        if self.location < self.rows.len() {
            Some(&self.rows[self.location])
        }
        else {
            None
        }
    }
}

pub enum RowIterator<'a> {
    Postgres(postgres::rows::Iter<'a>),
    Test(TestRowIterator<'a>)
}

impl<'a> Iterator for RowIterator<'a> {
    type Item = Row<'a>;

    fn next(&mut self) -> Option<Row<'a>> {
        match self {
            &mut RowIterator::Postgres(ref mut rows) => {
                rows.next().map(|r| Row::Postgres(r))
            }
            &mut RowIterator::Test(ref mut rows) => {
                rows.next().map(|r| Row::Test(r))
            }
        }
    }
}

pub enum Rows {
    Postgres(postgres::rows::Rows),
    Test(Vec<TestRow>)
}

impl<'stmt> Rows {
    pub fn get<'a>(&'a self, id: usize) -> Row<'a> {
        match self {
            &Rows::Postgres(ref rows) => {
                Row::Postgres(rows.get(id))
            }
            &Rows::Test(ref rows) => {
                Row::Test(rows.get(id).unwrap())
            }
        }
    }
    pub fn is_empty(&self) -> bool {
        match self {
            &Rows::Postgres(ref rows) => {
                rows.is_empty()
            }
            &Rows::Test(ref rows) => {
                rows.is_empty()
            }
        }
    }
    pub fn len(&self) -> usize {
        match self {
            &Rows::Postgres(ref rows) => {
                rows.len()
            }
            &Rows::Test(ref rows) => {
                rows.len()
            }
        }
    }
    pub fn iter<'a>(&'a self) -> RowIterator<'a> {
        match self {
            &Rows::Postgres(ref rows) => {
                RowIterator::Postgres(rows.iter())
            }
            &Rows::Test(ref rows) => {
                RowIterator::Test(TestRowIterator::new(rows))
            }
        }
    }
}

impl<'a> iter::IntoIterator for &'a Rows {
    type Item = Row<'a>;
    type IntoIter = RowIterator<'a>;
    fn into_iter(self) -> RowIterator<'a> {
        self.iter()
    }
}

#[derive(Debug, Clone)]
pub struct TestConnection {
    query_results: HashMap<String, Vec<TestRow>>,
    execute_results: HashMap<String, u64>
}

impl TestConnection {
    pub fn new() -> TestConnection {
        TestConnection {
            query_results: HashMap::new(),
            execute_results: HashMap::new()
        }
    }

    pub fn add_test_query<C>(&mut self, cmd: C, results: Vec<TestRow>)
        where C: Into<String> {
        self.query_results.insert(cmd.into(), results);
    }

    pub fn add_test_execute<C>(&mut self, cmd: C, results: u64)
        where C: Into<String> {
        self.execute_results.insert(cmd.into(), results);
    }

    fn get_rows(&self, cmd: &str) -> Vec<TestRow> {
        self.query_results.get(cmd).ok_or_else(|| ErrorKind::NoTestQuery(String::from(cmd))).unwrap().clone()
    }
    fn execute(&self, cmd: &str) -> Result<u64> {
        self.execute_results.get(cmd).map(|i| *i ).ok_or_else(|| Error::from(ErrorKind::NoTestExecute(String::from(cmd))))
    }
}

#[derive(Debug)]
pub enum Connection {
    Postgres(r2d2::PooledConnection<r2d2_postgres::PostgresConnectionManager>),
    Test(TestConnection)
}
impl<'conn> Connection {
    pub fn query(&'conn self, query: &str) -> Result<Rows> {
        match self {
            &Connection::Postgres(ref conn) => {
                Ok(Rows::Postgres(conn.query(query, &[])?))
            }
            &Connection::Test(ref conn) => {
                Ok(Rows::Test(conn.get_rows(query)))
            }
        }
    }
    pub fn execute(&self, equery: &str) -> Result<u64> {
        match self {
            &Connection::Postgres(ref conn) => {
                conn.execute(equery, &[]).map_err(|e|From::from(e))
            }
            &Connection::Test(ref conn) => {
                conn.execute(equery)
            }
        }
    }
}

#[derive(Clone, Debug)]
pub enum Pool {
    Postgres(r2d2::Pool<r2d2_postgres::PostgresConnectionManager>),
    TestPool(TestConnection)
}

impl Pool {
    pub fn get(&self) -> Result<Connection> {
        match self {
            &Pool::Postgres(ref pool) => {
                let conn = pool.get()?;
                Ok(Connection::Postgres(conn))
            }
            &Pool::TestPool(ref conn) => {
                Ok(Connection::Test(conn.clone()))
            }
        }
    }
}

pub struct PoolKey;

impl Key for PoolKey {
    type Value = Pool;
}

// Gets a connection from the pool from the given request or returns a 500
#[macro_export]
macro_rules! get_db_connection {
    ($req:expr) => (match $req.extensions.get::<persistent::Read<db::PoolKey>>() {
        Some(pool) => match pool.get() {
            Ok(conn) => {
                conn
            }
            Err(_) => {
                println!("Couldn't get a connection to pg!");
                return Ok(Response::with(status::InternalServerError));
            }
        },
        None => {
            println!("Couldn't get the pg pool from the request!");
            return Ok(Response::with(status::InternalServerError));
        }
    })
}
<commit_msg>If we get a Postgres error, also return the command<commit_after>use iron::typemap::Key;
use uuid::Uuid;
use serde_json;
use std::iter;
use r2d2;
use r2d2_postgres;
use postgres;
use std::convert::From;
use std::collections::HashMap;
use std::fmt;

error_chain! {
    errors {
        UniqueViolation
        NoTestQuery(cmd: String)
        NoTestExecute(cmd: String)
        PostgresError(cmd: String)
    }
    foreign_links {
        R2D2Error(r2d2::Error);
    }
}

pub struct HexSlice<'a>(&'a [u8]);

impl<'a> HexSlice<'a> {
    pub fn new<T>(data: &'a T) -> HexSlice<'a>
        where T: ?Sized + AsRef<[u8]> + 'a
    {
        HexSlice(data.as_ref())
    }
}

impl<'a> fmt::Display for HexSlice<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for byte in self.0 {
            write!(f, "{:02X}", byte)?;
        }
        Ok(())
    }
}

pub trait FromSql {}
impl FromSql for u32 {}
impl FromSql for Uuid {}
impl FromSql for String {}
impl <'a> FromSql for &'a str {}
impl FromSql for Vec<u8> {}
impl FromSql for serde_json::Value {}

#[derive(Debug, Clone)]
pub enum SqlValue {
    U32(u32),
    UUID(Uuid),
    String(String),
    U8Bytes(Vec<u8>),
    JSON(serde_json::Value),
}

impl From<u32> for SqlValue {
    fn from(i: u32) -> SqlValue { SqlValue::U32(i) }
}

#[derive(Hash, Eq, PartialEq, Clone, Debug)]
pub enum ValueIndex
{
    U32(u32),
    String(String),
}

impl<'a> From<&'a str> for ValueIndex {
    fn from(s: &str) -> ValueIndex { ValueIndex::String(String::from(s))}
}

pub trait RowIndex {
    fn val(&self) -> ValueIndex;
}
impl RowIndex for u32 {
    fn val(&self) -> ValueIndex { ValueIndex::U32(*self)}
}
impl RowIndex for String {
    fn val(&self) -> ValueIndex { ValueIndex::String(self.clone())}
}
impl <'a> RowIndex for &'a str {
    fn val(&self) -> ValueIndex { ValueIndex::String(String::from(*self))}
}

#[derive(Debug, Clone)]
pub struct TestRow
{
    data: HashMap<ValueIndex, SqlValue>
}

pub trait GetRow<T> {
    fn get<R>(&self, id: R) -> T where T: FromSql, R: RowIndex + fmt::Display;
}

impl TestRow {
    pub fn new() -> TestRow {
        TestRow{data: HashMap::new()}
    }

    pub fn insert<K, V>(&mut self, k: K, v: V) -> Option<SqlValue>
        where K: Into<ValueIndex>, V: Into<SqlValue> {
        self.data.insert(k.into(), v.into())
    }
}

macro_rules! get_row {
    ($type: ty, $kind: path) => (impl<'a> GetRow<$type> for TestRow {
    fn get<R>(&self, id: R) -> $type where R: RowIndex + fmt::Display {
        if !self.data.contains_key(&id.val()) {
            panic!(format!("Can't find key {} in row", id));
        }
        match self.data[&id.val()] {
            $kind(ref val) => val.clone(),
            _ => panic!()
        }
    }
})}

get_row!(u32, SqlValue::U32);
get_row!(Uuid, SqlValue::UUID);
get_row!(String, SqlValue::String);
get_row!(serde_json::Value, SqlValue::JSON);
get_row!(Vec<u8>, SqlValue::U8Bytes);

pub enum Row<'a> {
    Postgres(postgres::rows::Row<'a>),
    Test(&'a TestRow)
}
impl<'a> Row<'a> {
    pub fn get<T, R>(&self, id: R) -> T
    where
        T: FromSql + postgres::types::FromSql,
        R: RowIndex + postgres::rows::RowIndex + fmt::Display + fmt::Debug,
        TestRow: GetRow<T> {
        match self {
            &Row::Postgres(ref rows) => {
                rows.get(id)
            }
            &Row::Test(ref rows) => {
                rows.get(id)
            }
        }
    }
    pub fn get_opt<T, R>(&self, id: R) -> Option<Result<T>> where T: FromSql, R: RowIndex {
        unimplemented!();
    }
}

pub struct TestRowIterator<'a> {
    rows: &'a Vec<TestRow>,
    location: usize
}

impl<'a> TestRowIterator<'a> {
    fn new(r: &'a Vec<TestRow>) -> TestRowIterator<'a> {
        TestRowIterator {
            rows: r,
            location: 0
        }
    }

    fn next(&mut self) -> Option<&'a TestRow> {
        self.location +=1;
        if self.location < self.rows.len() {
            Some(&self.rows[self.location])
        }
        else {
            None
        }
    }
}

pub enum RowIterator<'a> {
    Postgres(postgres::rows::Iter<'a>),
    Test(TestRowIterator<'a>)
}

impl<'a> Iterator for RowIterator<'a> {
    type Item = Row<'a>;

    fn next(&mut self) -> Option<Row<'a>> {
        match self {
            &mut RowIterator::Postgres(ref mut rows) => {
                rows.next().map(|r| Row::Postgres(r))
            }
            &mut RowIterator::Test(ref mut rows) => {
                rows.next().map(|r| Row::Test(r))
            }
        }
    }
}

pub enum Rows {
    Postgres(postgres::rows::Rows),
    Test(Vec<TestRow>)
}

impl<'stmt> Rows {
    pub fn get<'a>(&'a self, id: usize) -> Row<'a> {
        match self {
            &Rows::Postgres(ref rows) => {
                Row::Postgres(rows.get(id))
            }
            &Rows::Test(ref rows) => {
                Row::Test(rows.get(id).unwrap())
            }
        }
    }
    pub fn is_empty(&self) -> bool {
        match self {
            &Rows::Postgres(ref rows) => {
                rows.is_empty()
            }
            &Rows::Test(ref rows) => {
                rows.is_empty()
            }
        }
    }
    pub fn len(&self) -> usize {
        match self {
            &Rows::Postgres(ref rows) => {
                rows.len()
            }
            &Rows::Test(ref rows) => {
                rows.len()
            }
        }
    }
    pub fn iter<'a>(&'a self) -> RowIterator<'a> {
        match self {
            &Rows::Postgres(ref rows) => {
                RowIterator::Postgres(rows.iter())
            }
            &Rows::Test(ref rows) => {
                RowIterator::Test(TestRowIterator::new(rows))
            }
        }
    }
}

impl<'a> iter::IntoIterator for &'a Rows {
    type Item = Row<'a>;
    type IntoIter = RowIterator<'a>;
    fn into_iter(self) -> RowIterator<'a> {
        self.iter()
    }
}

#[derive(Debug, Clone)]
pub struct TestConnection {
    query_results: HashMap<String, Vec<TestRow>>,
    execute_results: HashMap<String, u64>
}

impl TestConnection {
    pub fn new() -> TestConnection {
        TestConnection {
            query_results: HashMap::new(),
            execute_results: HashMap::new()
        }
    }

    pub fn add_test_query<C>(&mut self, cmd: C, results: Vec<TestRow>)
        where C: Into<String> {
        self.query_results.insert(cmd.into(), results);
    }

    pub fn add_test_execute<C>(&mut self, cmd: C, results: u64)
        where C: Into<String> {
        self.execute_results.insert(cmd.into(), results);
    }

    fn get_rows(&self, cmd: &str) -> Vec<TestRow> {
        self.query_results.get(cmd).ok_or_else(|| ErrorKind::NoTestQuery(String::from(cmd))).unwrap().clone()
    }
    fn execute(&self, cmd: &str) -> Result<u64> {
        self.execute_results.get(cmd).map(|i| *i ).ok_or_else(|| Error::from(ErrorKind::NoTestExecute(String::from(cmd))))
    }
}

#[derive(Debug)]
pub enum Connection {
    Postgres(r2d2::PooledConnection<r2d2_postgres::PostgresConnectionManager>),
    Test(TestConnection)
}
impl<'conn> Connection {
    pub fn query(&'conn self, query: &str) -> Result<Rows> {
        match self {
            &Connection::Postgres(ref conn) => {
                Ok(Rows::Postgres(conn.query(query, &[]).map_err(|e| Error::with_chain(e, ErrorKind::PostgresError(query.to_string())))?))
            }
            &Connection::Test(ref conn) => {
                Ok(Rows::Test(conn.get_rows(query)))
            }
        }
    }
    pub fn execute(&self, equery: &str) -> Result<u64> {
        match self {
            &Connection::Postgres(ref conn) => {
                conn.execute(equery, &[]).map_err(|e| Error::with_chain(e, ErrorKind::PostgresError(equery.to_string())))
            }
            &Connection::Test(ref conn) => {
                conn.execute(equery)
            }
        }
    }
}

#[derive(Clone, Debug)]
pub enum Pool {
    Postgres(r2d2::Pool<r2d2_postgres::PostgresConnectionManager>),
    TestPool(TestConnection)
}

impl Pool {
    pub fn get(&self) -> Result<Connection> {
        match self {
            &Pool::Postgres(ref pool) => {
                let conn = pool.get()?;
                Ok(Connection::Postgres(conn))
            }
            &Pool::TestPool(ref conn) => {
                Ok(Connection::Test(conn.clone()))
            }
        }
    }
}

pub struct PoolKey;

impl Key for PoolKey {
    type Value = Pool;
}

// Gets a connection from the pool from the given request or returns a 500
#[macro_export]
macro_rules! get_db_connection {
    ($req:expr) => (match $req.extensions.get::<persistent::Read<db::PoolKey>>() {
        Some(pool) => match pool.get() {
            Ok(conn) => {
                conn
            }
            Err(_) => {
                println!("Couldn't get a connection to pg!");
                return Ok(Response::with(status::InternalServerError));
            }
        },
        None => {
            println!("Couldn't get the pg pool from the request!");
            return Ok(Response::with(status::InternalServerError));
        }
    })
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::io::Write;
use std::str::FromStr;

use bufstream::BufStream;

use protocol::cmd::Cmd;
use protocol::cmd::Get;
use protocol::cmd::Resp;
use protocol::cmd::Set;

use super::errors::TcpTransportError;
use super::stats::TransportStats;
use super::typedefs::TcpTransportResult;

pub struct TcpTransport<T: Read + Write> {
    stream: BufStream<T>,

    metrics: TransportStats,
    key_maxlen: u64,
}

impl<T: Read + Write> TcpTransport<T> {
    pub fn new(stream: T) -> TcpTransport<T> {
        TcpTransport {
            key_maxlen: 250, // memcached standard
            metrics: TransportStats::new(),
            stream: BufStream::new(stream),
        }
    }

    pub fn with_key_maxlen(&mut self,
                           key_maxlen: u64)
                           -> &mut TcpTransport<T> {
        self.key_maxlen = key_maxlen;
        self
    }

    pub fn get_max_line_len(&self) -> usize {
        // This needs to be the length of the longest command line, not
        // including data values for which the length is given upfront
        self.key_maxlen as usize + 100
    }

    pub fn get_stats_clone(&self) -> TransportStats {
        self.metrics.clone()
    }

    pub fn get_stream(&self) -> &T {
        self.stream.get_ref()
    }

    // Basic bytes manipulation and reading from the stream

    pub fn as_string(&self, bytes: Vec<u8>) -> TcpTransportResult<String> {
        match String::from_utf8(bytes) {
            Ok(string) => Ok(string),
            Err(_) => Err(TcpTransportError::Utf8Error),
        }
    }

    pub fn as_number<N: FromStr>(&self,
                                 bytes: Vec<u8>)
                                 -> TcpTransportResult<N> {
        let string = try!(self.as_string(bytes));
        match string.parse::<N>() {
            Ok(num) => Ok(num),
            Err(_) => Err(TcpTransportError::NumberParseError),
        }
    }

    pub fn read_byte(&mut self) -> TcpTransportResult<u8> {
        let mut bytes = [0; 1];

        match self.stream.read(&mut bytes) {
            Ok(1) => {
                // Update metrics
                self.metrics.bytes_read += 1;

                Ok(bytes[0])
            }
            _ => Err(TcpTransportError::StreamReadError),
        }
    }

    pub fn read_bytes(&mut self, len: u64) -> TcpTransportResult<Vec<u8>> {
        let mut bytes = vec![];

        for _ in 0..len {
            let byte = try!(self.read_byte());
            bytes.push(byte);
        }

        Ok(bytes)
    }

    pub fn read_line(&mut self, maxlen: usize) -> TcpTransportResult<Vec<u8>> {
        let mut bytes = vec![];
        let mut found_line_end = false;

        for _ in 0..maxlen {
            let byte = try!(self.read_byte());
            bytes.push(byte);

            // Look for \r\n
            if bytes.ends_with(&[13, 10]) {
                found_line_end = true;
                break;
            }
        }

        if found_line_end {
            // Pop off \r\n
            bytes.pop();
            bytes.pop();
            Ok(bytes)
        } else {
            Err(TcpTransportError::LineReadError)
        }
    }

    pub fn remove_first_char(&self,
                             bytes: &mut Vec<u8>)
                             -> TcpTransportResult<()> {
        match bytes.len() > 0 {
            true => {
                bytes.remove(0);
                Ok(())
            }
            false => Err(TcpTransportError::StreamReadError),
        }
    }

    pub fn parse_word(&self,
                      bytes: Vec<u8>)
                      -> TcpTransportResult<(Vec<u8>, Vec<u8>)> {
        let mut space_idx: i64 = -1;

        for i in 0..bytes.len() {
            // We're looking for a space
            if bytes[i] == 32 {
                space_idx = i as i64;
                break;
            }
        }

        if space_idx > -1 {
            let mut word = vec![];
            let mut rest = vec![];

            // TODO figure out how to return a modified vector instead of
            // copying the whole rest of it
            for i in 0..bytes.len() {
                let byte = bytes[i];
                if (i as i64) < space_idx {
                    word.push(byte);
                } else {
                    rest.push(byte);
                }
            }

            Ok((word, rest))

        } else {
            // If we've reached the end of the buffer without seeing a space
            // that makes the whole buffer a word
            Ok((bytes, vec![]))
        }
    }

    // Writing to the stream

    pub fn flush_writes(&mut self) -> TcpTransportResult<()> {
        match self.stream.flush() {
            Ok(_) => Ok(()),
            Err(_) => Err(TcpTransportError::StreamWriteError),
        }
    }

    pub fn write_bytes(&mut self,
                       bytes: &Vec<u8>)
                       -> TcpTransportResult<usize> {
        match self.stream.write(bytes) {
            Ok(cnt_written) => {
                // Update metrics
                self.metrics.bytes_written += cnt_written as u64;

                Ok(cnt_written)
            }
            Err(_) => Err(TcpTransportError::StreamWriteError)
        }
    }

    pub fn write_string(&mut self, string: &str) -> TcpTransportResult<usize> {
        let bytes = string.to_string().into_bytes();
        Ok(try!(self.write_bytes(&bytes)))
    }

    // Parse individual commands

    pub fn parse_cmd_get(&self, mut rest: Vec<u8>) -> TcpTransportResult<Cmd> {
        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (key, rest) = try!(self.parse_word(rest));

        // We expect to find the end of the line now
        if rest.is_empty() {
            let key_str = try!(self.as_string(key));
            Ok(Cmd::Get(Get { key: key_str }))
        } else {
            Err(TcpTransportError::CommandParseError)
        }
    }

    pub fn parse_cmd_set(&mut self,
                         mut rest: Vec<u8>)
                         -> TcpTransportResult<Cmd> {
        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (key, mut rest) = try!(self.parse_word(rest));

        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (flags, mut rest) = try!(self.parse_word(rest));

        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (exptime, mut rest) = try!(self.parse_word(rest));

        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (bytelen, _) = try!(self.parse_word(rest));

        let key_str = try!(self.as_string(key));
        let flags_num = try!(self.as_number::<u16>(flags));
        let exptime_num = try!(self.as_number::<u32>(exptime));
        let bytelen_num = try!(self.as_number::<u64>(bytelen));

        // We know the byte length, so now read the value
        let value = try!(self.read_bytes(bytelen_num));

        // Read the line termination marker
        let line_len = self.get_max_line_len();
        let rest = try!(self.read_line(line_len));

        // We got all the values we expected and there is nothing left
        if rest.is_empty() {
            return Ok(Cmd::Set(Set {
                key: key_str,
                exptime: exptime_num,
                data: value,
            }));
        }

        Err(TcpTransportError::CommandParseError)
    }

    // High level functions

    pub fn read_cmd(&mut self) -> TcpTransportResult<Cmd> {
        let line_len = self.get_max_line_len();

        let fst_line = try!(self.read_line(line_len));
        let (keyword, rest) = try!(self.parse_word(fst_line));
        let keyword_str = try!(self.as_string(keyword));

        if keyword_str == "get" {
            return self.parse_cmd_get(rest);
        } else if keyword_str == "set" {
            return self.parse_cmd_set(rest);
        } else if keyword_str == "stats" {
            return Ok(Cmd::Stats);
        }

        Err(TcpTransportError::InvalidCmd)
    }

    pub fn write_resp(&mut self, resp: &Resp) -> TcpTransportResult<()> {
        match *resp {
            Resp::Error => {
                try!(self.write_string("ERROR\r\n"));
            }
            Resp::Stats(ref stats) => {
                for stat in stats {
                    try!(self.write_string("STAT "));
                    try!(self.write_string(&stat.key));
                    try!(self.write_string(" "));
                    try!(self.write_string(&stat.value));
                    try!(self.write_string("\r\n"));
                }
                try!(self.write_string("END\r\n"));
            }
            Resp::Stored => {
                try!(self.write_string("STORED\r\n"));
            }
            Resp::Value(ref value) => {
                try!(self.write_string("VALUE ")); // keyword
                try!(self.write_string(&value.key)); // key
                try!(self.write_string(" 0 ")); // flags
                try!(self.write_string(&value.data.len().to_string())); // bytelen
                try!(self.write_string(&"\r\n".to_string())); // newline
                try!(self.write_bytes(&value.data)); // data block
                try!(self.write_string(&"\r\n".to_string())); // newline
            }
            _ => {
                return Err(TcpTransportError::StreamWriteError);
            }
        }

        // Make sure all bytes were actually sent
        self.flush_writes()
    }
}
<commit_msg>fix mentions of metrics in transport that are supposed to be stats<commit_after>use std::io::Read;
use std::io::Write;
use std::str::FromStr;

use bufstream::BufStream;

use protocol::cmd::Cmd;
use protocol::cmd::Get;
use protocol::cmd::Resp;
use protocol::cmd::Set;

use super::errors::TcpTransportError;
use super::stats::TransportStats;
use super::typedefs::TcpTransportResult;

pub struct TcpTransport<T: Read + Write> {
    stream: BufStream<T>,

    stats: TransportStats,
    key_maxlen: u64,
}

impl<T: Read + Write> TcpTransport<T> {
    pub fn new(stream: T) -> TcpTransport<T> {
        TcpTransport {
            key_maxlen: 250, // memcached standard
            stats: TransportStats::new(),
            stream: BufStream::new(stream),
        }
    }

    pub fn with_key_maxlen(&mut self,
                           key_maxlen: u64)
                           -> &mut TcpTransport<T> {
        self.key_maxlen = key_maxlen;
        self
    }

    pub fn get_max_line_len(&self) -> usize {
        // This needs to be the length of the longest command line, not
        // including data values for which the length is given upfront
        self.key_maxlen as usize + 100
    }

    pub fn get_stats_clone(&self) -> TransportStats {
        self.stats.clone()
    }

    pub fn get_stream(&self) -> &T {
        self.stream.get_ref()
    }

    // Basic bytes manipulation and reading from the stream

    pub fn as_string(&self, bytes: Vec<u8>) -> TcpTransportResult<String> {
        match String::from_utf8(bytes) {
            Ok(string) => Ok(string),
            Err(_) => Err(TcpTransportError::Utf8Error),
        }
    }

    pub fn as_number<N: FromStr>(&self,
                                 bytes: Vec<u8>)
                                 -> TcpTransportResult<N> {
        let string = try!(self.as_string(bytes));
        match string.parse::<N>() {
            Ok(num) => Ok(num),
            Err(_) => Err(TcpTransportError::NumberParseError),
        }
    }

    pub fn read_byte(&mut self) -> TcpTransportResult<u8> {
        let mut bytes = [0; 1];

        match self.stream.read(&mut bytes) {
            Ok(1) => {
                // Update stats
                self.stats.bytes_read += 1;

                Ok(bytes[0])
            }
            _ => Err(TcpTransportError::StreamReadError),
        }
    }

    pub fn read_bytes(&mut self, len: u64) -> TcpTransportResult<Vec<u8>> {
        let mut bytes = vec![];

        for _ in 0..len {
            let byte = try!(self.read_byte());
            bytes.push(byte);
        }

        Ok(bytes)
    }

    pub fn read_line(&mut self, maxlen: usize) -> TcpTransportResult<Vec<u8>> {
        let mut bytes = vec![];
        let mut found_line_end = false;

        for _ in 0..maxlen {
            let byte = try!(self.read_byte());
            bytes.push(byte);

            // Look for \r\n
            if bytes.ends_with(&[13, 10]) {
                found_line_end = true;
                break;
            }
        }

        if found_line_end {
            // Pop off \r\n
            bytes.pop();
            bytes.pop();
            Ok(bytes)
        } else {
            Err(TcpTransportError::LineReadError)
        }
    }

    pub fn remove_first_char(&self,
                             bytes: &mut Vec<u8>)
                             -> TcpTransportResult<()> {
        match bytes.len() > 0 {
            true => {
                bytes.remove(0);
                Ok(())
            }
            false => Err(TcpTransportError::StreamReadError),
        }
    }

    pub fn parse_word(&self,
                      bytes: Vec<u8>)
                      -> TcpTransportResult<(Vec<u8>, Vec<u8>)> {
        let mut space_idx: i64 = -1;

        for i in 0..bytes.len() {
            // We're looking for a space
            if bytes[i] == 32 {
                space_idx = i as i64;
                break;
            }
        }

        if space_idx > -1 {
            let mut word = vec![];
            let mut rest = vec![];

            // TODO figure out how to return a modified vector instead of
            // copying the whole rest of it
            for i in 0..bytes.len() {
                let byte = bytes[i];
                if (i as i64) < space_idx {
                    word.push(byte);
                } else {
                    rest.push(byte);
                }
            }

            Ok((word, rest))

        } else {
            // If we've reached the end of the buffer without seeing a space
            // that makes the whole buffer a word
            Ok((bytes, vec![]))
        }
    }

    // Writing to the stream

    pub fn flush_writes(&mut self) -> TcpTransportResult<()> {
        match self.stream.flush() {
            Ok(_) => Ok(()),
            Err(_) => Err(TcpTransportError::StreamWriteError),
        }
    }

    pub fn write_bytes(&mut self,
                       bytes: &Vec<u8>)
                       -> TcpTransportResult<usize> {
        match self.stream.write(bytes) {
            Ok(cnt_written) => {
                // Update stats
                self.stats.bytes_written += cnt_written as u64;

                Ok(cnt_written)
            }
            Err(_) => Err(TcpTransportError::StreamWriteError)
        }
    }

    pub fn write_string(&mut self, string: &str) -> TcpTransportResult<usize> {
        let bytes = string.to_string().into_bytes();
        Ok(try!(self.write_bytes(&bytes)))
    }

    // Parse individual commands

    pub fn parse_cmd_get(&self, mut rest: Vec<u8>) -> TcpTransportResult<Cmd> {
        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (key, rest) = try!(self.parse_word(rest));

        // We expect to find the end of the line now
        if rest.is_empty() {
            let key_str = try!(self.as_string(key));
            Ok(Cmd::Get(Get { key: key_str }))
        } else {
            Err(TcpTransportError::CommandParseError)
        }
    }

    pub fn parse_cmd_set(&mut self,
                         mut rest: Vec<u8>)
                         -> TcpTransportResult<Cmd> {
        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (key, mut rest) = try!(self.parse_word(rest));

        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (flags, mut rest) = try!(self.parse_word(rest));

        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (exptime, mut rest) = try!(self.parse_word(rest));

        try!(self.remove_first_char(&mut rest)); // remove leading space
        let (bytelen, _) = try!(self.parse_word(rest));

        let key_str = try!(self.as_string(key));
        let flags_num = try!(self.as_number::<u16>(flags));
        let exptime_num = try!(self.as_number::<u32>(exptime));
        let bytelen_num = try!(self.as_number::<u64>(bytelen));

        // We know the byte length, so now read the value
        let value = try!(self.read_bytes(bytelen_num));

        // Read the line termination marker
        let line_len = self.get_max_line_len();
        let rest = try!(self.read_line(line_len));

        // We got all the values we expected and there is nothing left
        if rest.is_empty() {
            return Ok(Cmd::Set(Set {
                key: key_str,
                exptime: exptime_num,
                data: value,
            }));
        }

        Err(TcpTransportError::CommandParseError)
    }

    // High level functions

    pub fn read_cmd(&mut self) -> TcpTransportResult<Cmd> {
        let line_len = self.get_max_line_len();

        let fst_line = try!(self.read_line(line_len));
        let (keyword, rest) = try!(self.parse_word(fst_line));
        let keyword_str = try!(self.as_string(keyword));

        if keyword_str == "get" {
            return self.parse_cmd_get(rest);
        } else if keyword_str == "set" {
            return self.parse_cmd_set(rest);
        } else if keyword_str == "stats" {
            return Ok(Cmd::Stats);
        }

        Err(TcpTransportError::InvalidCmd)
    }

    pub fn write_resp(&mut self, resp: &Resp) -> TcpTransportResult<()> {
        match *resp {
            Resp::Error => {
                try!(self.write_string("ERROR\r\n"));
            }
            Resp::Stats(ref stats) => {
                for stat in stats {
                    try!(self.write_string("STAT "));
                    try!(self.write_string(&stat.key));
                    try!(self.write_string(" "));
                    try!(self.write_string(&stat.value));
                    try!(self.write_string("\r\n"));
                }
                try!(self.write_string("END\r\n"));
            }
            Resp::Stored => {
                try!(self.write_string("STORED\r\n"));
            }
            Resp::Value(ref value) => {
                try!(self.write_string("VALUE ")); // keyword
                try!(self.write_string(&value.key)); // key
                try!(self.write_string(" 0 ")); // flags
                try!(self.write_string(&value.data.len().to_string())); // bytelen
                try!(self.write_string(&"\r\n".to_string())); // newline
                try!(self.write_bytes(&value.data)); // data block
                try!(self.write_string(&"\r\n".to_string())); // newline
            }
            _ => {
                return Err(TcpTransportError::StreamWriteError);
            }
        }

        // Make sure all bytes were actually sent
        self.flush_writes()
    }
}
<|endoftext|>


<commit_before>use crate::{
    crypto::{double_sha256, Digest, KeyPair, ScriptHash, SigPair},
    serializer::*,
    tx::*,
};
use std::{collections::BTreeSet, io::Cursor, ops::Deref, sync::Arc};

pub type BlockFilter = BTreeSet<ScriptHash>;

#[derive(Clone, Debug, PartialEq)]
pub enum FilteredBlock {
    Header((BlockHeader, SigPair)),
    Block(Arc<Block>),
}

#[derive(Clone, Debug, PartialEq)]
pub enum Block {
    V0(BlockV0),
}

impl Block {
    #[inline]
    pub fn header(&self) -> BlockHeader {
        match self {
            Block::V0(block) => BlockHeader::V0(block.header.clone()),
        }
    }

    #[inline]
    pub fn height(&self) -> u64 {
        match self {
            Block::V0(block) => block.height,
        }
    }

    #[inline]
    pub fn timestamp(&self) -> u64 {
        match self {
            Block::V0(block) => block.timestamp,
        }
    }

    #[inline]
    pub fn txs(&self) -> &[TxVariant] {
        match self {
            Block::V0(block) => &block.transactions,
        }
    }

    #[inline]
    pub fn signer(&self) -> Option<&SigPair> {
        match self {
            Block::V0(block) => block.signer.as_ref(),
        }
    }

    pub fn sign(&mut self, key_pair: &KeyPair) {
        let buf = self.calc_header_hash();
        match self {
            Block::V0(block) => {
                block.signer = Some(key_pair.sign(buf.as_ref()));
            }
        }
    }

    pub fn verify_previous_hash(&self, prev_block: &Self) -> bool {
        let cur_prev_hash = match self {
            Block::V0(block) => &block.previous_hash,
        };
        cur_prev_hash == &prev_block.calc_header_hash()
    }

    pub fn verify_tx_merkle_root(&self) -> bool {
        match self {
            Block::V0(block) => {
                let digest = calc_tx_merkle_root(&block.transactions);
                block.tx_merkle_root == digest
            }
        }
    }

    pub fn calc_header_hash(&self) -> Digest {
        match self {
            Block::V0(block) => block.calc_header_hash(),
        }
    }

    pub fn deserialize(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        let header = BlockHeader::deserialize(cur)?;
        match header {
            BlockHeader::V0(header) => {
                let signer = Some(cur.take_sig_pair().ok()?);

                let len = cur.take_u32().ok()?;
                let mut transactions = Vec::<TxVariant>::with_capacity(len as usize);
                for _ in 0..len {
                    transactions.push(TxVariant::deserialize(cur)?);
                }

                Some(Block::V0(BlockV0 {
                    header,
                    signer,
                    transactions,
                }))
            }
        }
    }

    pub fn serialize(&self, buf: &mut Vec<u8>) {
        match self {
            Block::V0(block) => {
                block.header.serialize(buf);
                buf.push_sig_pair(
                    block
                        .signer
                        .as_ref()
                        .expect("block must be signed to serialize"),
                );
                buf.push_u32(block.transactions.len() as u32);
                for tx in &block.transactions {
                    tx.serialize(buf)
                }
            }
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum BlockHeader {
    V0(BlockHeaderV0),
}

impl BlockHeader {
    pub fn serialize(&self, buf: &mut Vec<u8>) {
        match self {
            BlockHeader::V0(header) => header.serialize(buf),
        }
    }

    pub fn deserialize(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        let header_ver = cur.take_u16().ok()?;
        match header_ver {
            0x00 => Some(BlockHeader::V0(BlockHeaderV0::deserialize(cur)?)),
            _ => None,
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct BlockHeaderV0 {
    pub previous_hash: Digest,
    pub height: u64,
    pub timestamp: u64,
    pub tx_merkle_root: Digest,
}

impl BlockHeaderV0 {
    pub(self) fn serialize(&self, buf: &mut Vec<u8>) {
        // Header version (2 bytes)
        buf.push_u16(0x00);

        buf.push_digest(&self.previous_hash);
        buf.push_u64(self.height);
        buf.push_u64(self.timestamp);
        buf.push_digest(&self.tx_merkle_root);
    }

    pub(self) fn deserialize(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        // We expect the version to already be deserialized here

        let previous_hash = cur.take_digest().ok()?;
        let height = cur.take_u64().ok()?;
        let timestamp = cur.take_u64().ok()?;
        let tx_merkle_root = cur.take_digest().ok()?;
        Some(Self {
            previous_hash,
            height,
            timestamp,
            tx_merkle_root,
        })
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct BlockV0 {
    pub header: BlockHeaderV0,
    pub signer: Option<SigPair>,
    pub transactions: Vec<TxVariant>,
}

impl BlockV0 {
    pub fn new_child(&self, txs: Vec<TxVariant>) -> Block {
        let previous_hash = self.calc_header_hash();
        let height = self.header.height + 1;
        let tx_merkle_root = calc_tx_merkle_root(&txs);
        let timestamp = crate::get_epoch_time();
        Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash,
                height,
                timestamp,
                tx_merkle_root,
            },
            signer: None,
            transactions: txs,
        })
    }

    pub fn calc_header_hash(&self) -> Digest {
        let mut buf = Vec::with_capacity(1024);
        self.header.serialize(&mut buf);
        double_sha256(&buf)
    }
}

impl Deref for BlockV0 {
    type Target = BlockHeaderV0;

    fn deref(&self) -> &BlockHeaderV0 {
        &self.header
    }
}

pub fn calc_tx_merkle_root(txs: &[TxVariant]) -> Digest {
    let mut buf = Vec::with_capacity(4096 * txs.len());
    for tx in txs {
        tx.serialize(&mut buf)
    }
    double_sha256(&buf)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{asset::Asset, crypto::KeyPair};

    #[test]
    fn serialize_block_v0() {
        let keys = KeyPair::gen();
        let transactions = vec![TxVariant::V0(TxVariantV0::RewardTx(RewardTx {
            base: Tx {
                nonce: 111,
                expiry: 1234567890,
                fee: Asset::default(),
                signature_pairs: Vec::new(),
            },
            to: keys.0.clone().into(),
            rewards: Asset::default(),
        }))];
        let tx_merkle_root = {
            let mut buf = Vec::new();
            for tx in &transactions {
                tx.serialize(&mut buf)
            }
            double_sha256(&buf)
        };
        let mut block = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: Digest::from_slice(&[0u8; 32]).unwrap(),
                height: 123,
                timestamp: 1532992800,
                tx_merkle_root,
            },
            signer: None,
            transactions,
        });
        block.sign(&keys);

        let mut buf = Vec::new();
        block.serialize(&mut buf);

        let mut cur = Cursor::<&[u8]>::new(&buf);
        let dec = Block::deserialize(&mut cur).unwrap();

        assert_eq!(block, dec);
    }

    #[test]
    fn merkle_root() {
        let mut block = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: Digest::from_slice(&[0; 32]).unwrap(),
                height: 0,
                timestamp: 0,
                tx_merkle_root: double_sha256(&[0; 0]),
            },
            signer: None,
            transactions: vec![],
        });
        assert!(block.verify_tx_merkle_root());

        match &mut block {
            Block::V0(block) => {
                block.header.tx_merkle_root = Digest::from_slice(&[0; 32]).unwrap();
            }
        }

        assert!(!block.verify_tx_merkle_root());
    }

    #[test]
    fn previous_hash() {
        let block_0 = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: Digest::from_slice(&[0; 32]).unwrap(),
                height: 0,
                timestamp: 0,
                tx_merkle_root: double_sha256(&[0; 0]),
            },
            signer: None,
            transactions: vec![],
        });

        let block_1 = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: block_0.calc_header_hash(),
                height: 1,
                timestamp: 0,
                tx_merkle_root: double_sha256(&[0; 0]),
            },
            signer: None,
            transactions: vec![],
        });

        let block_1_invalid = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: Digest::from_slice(&[0; 32]).unwrap(),
                height: 1,
                timestamp: 0,
                tx_merkle_root: double_sha256(&[0; 0]),
            },
            signer: None,
            transactions: vec![],
        });

        assert!(block_1.verify_previous_hash(&block_0));
        assert!(!block_1_invalid.verify_previous_hash(&block_0));
    }
}
<commit_msg>Use streaming hash when calculating the tx merkle root<commit_after>use crate::{
    crypto::{double_sha256, Digest, DoubleSha256, KeyPair, ScriptHash, SigPair},
    serializer::*,
    tx::*,
};
use std::{collections::BTreeSet, io::Cursor, ops::Deref, sync::Arc};

pub type BlockFilter = BTreeSet<ScriptHash>;

#[derive(Clone, Debug, PartialEq)]
pub enum FilteredBlock {
    Header((BlockHeader, SigPair)),
    Block(Arc<Block>),
}

#[derive(Clone, Debug, PartialEq)]
pub enum Block {
    V0(BlockV0),
}

impl Block {
    #[inline]
    pub fn header(&self) -> BlockHeader {
        match self {
            Block::V0(block) => BlockHeader::V0(block.header.clone()),
        }
    }

    #[inline]
    pub fn height(&self) -> u64 {
        match self {
            Block::V0(block) => block.height,
        }
    }

    #[inline]
    pub fn timestamp(&self) -> u64 {
        match self {
            Block::V0(block) => block.timestamp,
        }
    }

    #[inline]
    pub fn txs(&self) -> &[TxVariant] {
        match self {
            Block::V0(block) => &block.transactions,
        }
    }

    #[inline]
    pub fn signer(&self) -> Option<&SigPair> {
        match self {
            Block::V0(block) => block.signer.as_ref(),
        }
    }

    pub fn sign(&mut self, key_pair: &KeyPair) {
        let buf = self.calc_header_hash();
        match self {
            Block::V0(block) => {
                block.signer = Some(key_pair.sign(buf.as_ref()));
            }
        }
    }

    pub fn verify_previous_hash(&self, prev_block: &Self) -> bool {
        let cur_prev_hash = match self {
            Block::V0(block) => &block.previous_hash,
        };
        cur_prev_hash == &prev_block.calc_header_hash()
    }

    pub fn verify_tx_merkle_root(&self) -> bool {
        match self {
            Block::V0(block) => {
                let digest = calc_tx_merkle_root(&block.transactions);
                block.tx_merkle_root == digest
            }
        }
    }

    pub fn calc_header_hash(&self) -> Digest {
        match self {
            Block::V0(block) => block.calc_header_hash(),
        }
    }

    pub fn deserialize(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        let header = BlockHeader::deserialize(cur)?;
        match header {
            BlockHeader::V0(header) => {
                let signer = Some(cur.take_sig_pair().ok()?);

                let len = cur.take_u32().ok()?;
                let mut transactions = Vec::<TxVariant>::with_capacity(len as usize);
                for _ in 0..len {
                    transactions.push(TxVariant::deserialize(cur)?);
                }

                Some(Block::V0(BlockV0 {
                    header,
                    signer,
                    transactions,
                }))
            }
        }
    }

    pub fn serialize(&self, buf: &mut Vec<u8>) {
        match self {
            Block::V0(block) => {
                block.header.serialize(buf);
                buf.push_sig_pair(
                    block
                        .signer
                        .as_ref()
                        .expect("block must be signed to serialize"),
                );
                buf.push_u32(block.transactions.len() as u32);
                for tx in &block.transactions {
                    tx.serialize(buf)
                }
            }
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum BlockHeader {
    V0(BlockHeaderV0),
}

impl BlockHeader {
    pub fn serialize(&self, buf: &mut Vec<u8>) {
        match self {
            BlockHeader::V0(header) => header.serialize(buf),
        }
    }

    pub fn deserialize(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        let header_ver = cur.take_u16().ok()?;
        match header_ver {
            0x00 => Some(BlockHeader::V0(BlockHeaderV0::deserialize(cur)?)),
            _ => None,
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct BlockHeaderV0 {
    pub previous_hash: Digest,
    pub height: u64,
    pub timestamp: u64,
    pub tx_merkle_root: Digest,
}

impl BlockHeaderV0 {
    pub(self) fn serialize(&self, buf: &mut Vec<u8>) {
        // Header version (2 bytes)
        buf.push_u16(0x00);

        buf.push_digest(&self.previous_hash);
        buf.push_u64(self.height);
        buf.push_u64(self.timestamp);
        buf.push_digest(&self.tx_merkle_root);
    }

    pub(self) fn deserialize(cur: &mut Cursor<&[u8]>) -> Option<Self> {
        // We expect the version to already be deserialized here

        let previous_hash = cur.take_digest().ok()?;
        let height = cur.take_u64().ok()?;
        let timestamp = cur.take_u64().ok()?;
        let tx_merkle_root = cur.take_digest().ok()?;
        Some(Self {
            previous_hash,
            height,
            timestamp,
            tx_merkle_root,
        })
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct BlockV0 {
    pub header: BlockHeaderV0,
    pub signer: Option<SigPair>,
    pub transactions: Vec<TxVariant>,
}

impl BlockV0 {
    pub fn new_child(&self, txs: Vec<TxVariant>) -> Block {
        let previous_hash = self.calc_header_hash();
        let height = self.header.height + 1;
        let tx_merkle_root = calc_tx_merkle_root(&txs);
        let timestamp = crate::get_epoch_time();
        Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash,
                height,
                timestamp,
                tx_merkle_root,
            },
            signer: None,
            transactions: txs,
        })
    }

    pub fn calc_header_hash(&self) -> Digest {
        let mut buf = Vec::with_capacity(1024);
        self.header.serialize(&mut buf);
        double_sha256(&buf)
    }
}

impl Deref for BlockV0 {
    type Target = BlockHeaderV0;

    fn deref(&self) -> &BlockHeaderV0 {
        &self.header
    }
}

pub fn calc_tx_merkle_root(txs: &[TxVariant]) -> Digest {
    let mut hasher = DoubleSha256::new();
    let mut buf = Vec::with_capacity(4096);
    for tx in txs {
        buf.clear();
        tx.serialize(&mut buf);
        hasher.update(&buf);
    }
    hasher.finalize()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{asset::Asset, crypto::KeyPair};

    #[test]
    fn serialize_block_v0() {
        let keys = KeyPair::gen();
        let transactions = vec![TxVariant::V0(TxVariantV0::RewardTx(RewardTx {
            base: Tx {
                nonce: 111,
                expiry: 1234567890,
                fee: Asset::default(),
                signature_pairs: Vec::new(),
            },
            to: keys.0.clone().into(),
            rewards: Asset::default(),
        }))];
        let tx_merkle_root = calc_tx_merkle_root(&transactions);
        let mut block = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: Digest::from_slice(&[0u8; 32]).unwrap(),
                height: 123,
                timestamp: 1532992800,
                tx_merkle_root,
            },
            signer: None,
            transactions,
        });
        block.sign(&keys);

        let mut buf = Vec::new();
        block.serialize(&mut buf);

        let mut cur = Cursor::<&[u8]>::new(&buf);
        let dec = Block::deserialize(&mut cur).unwrap();

        assert_eq!(block, dec);
    }

    #[test]
    fn merkle_root() {
        let mut block = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: Digest::from_slice(&[0; 32]).unwrap(),
                height: 0,
                timestamp: 0,
                tx_merkle_root: double_sha256(&[0; 0]),
            },
            signer: None,
            transactions: vec![],
        });
        assert!(block.verify_tx_merkle_root());

        match &mut block {
            Block::V0(block) => {
                block.header.tx_merkle_root = Digest::from_slice(&[0; 32]).unwrap();
            }
        }

        assert!(!block.verify_tx_merkle_root());
    }

    #[test]
    fn previous_hash() {
        let block_0 = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: Digest::from_slice(&[0; 32]).unwrap(),
                height: 0,
                timestamp: 0,
                tx_merkle_root: double_sha256(&[0; 0]),
            },
            signer: None,
            transactions: vec![],
        });

        let block_1 = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: block_0.calc_header_hash(),
                height: 1,
                timestamp: 0,
                tx_merkle_root: double_sha256(&[0; 0]),
            },
            signer: None,
            transactions: vec![],
        });

        let block_1_invalid = Block::V0(BlockV0 {
            header: BlockHeaderV0 {
                previous_hash: Digest::from_slice(&[0; 32]).unwrap(),
                height: 1,
                timestamp: 0,
                tx_merkle_root: double_sha256(&[0; 0]),
            },
            signer: None,
            transactions: vec![],
        });

        assert!(block_1.verify_previous_hash(&block_0));
        assert!(!block_1_invalid.verify_previous_hash(&block_0));
    }
}
<|endoftext|>