<commit_before>use std::fmt::{Display, Formatter, Result};
use num_cpus;
use types::{CoordSum, Depth, Time};

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum UcbType {
  Winrate,
  Ucb1,
  Ucb1Tuned
}

impl Display for UcbType {
  fn fmt(&self, f: &mut Formatter) -> Result {
    match self {
      &UcbType::Winrate => write!(f, "Winrate"),
      &UcbType::Ucb1 => write!(f, "Ucb1"),
      &UcbType::Ucb1Tuned => write!(f, "Ucb1Tuned")
    }
  }
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum UctKomiType {
  None,
  Static,
  Dynamic
}

impl Display for UctKomiType {
  fn fmt(&self, f: &mut Formatter) -> Result {
    match self {
      &UctKomiType::None => write!(f, "None"),
      &UctKomiType::Static => write!(f, "Static"),
      &UctKomiType::Dynamic => write!(f, "Dynamic")
    }
  }
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum Solver {
  Uct,
  Heuristic
}

impl Display for Solver {
  fn fmt(&self, f: &mut Formatter) -> Result {
    match self {
      &Solver::Uct => write!(f, "Uct"),
      &Solver::Heuristic => write!(f, "Heuristic")
    }
  }
}

static UCT_RADIUS: CoordSum = 3;

static UCB_TYPE: UcbType = UcbType::Ucb1Tuned;

static FINAL_UCB_TYPE: UcbType = UcbType::Winrate;

static UCT_DRAW_WEIGHT: f32 = 0.4;

static UCTK: f32 = 1.0;

static UCT_WHEN_CREATE_CHILDREN: usize = 2;

static UCT_DEPTH: Depth = 8;

static mut THREADS_COUNT: usize = 4;

static UCT_KOMI_TYPE: UctKomiType = UctKomiType::Static;

static UCT_RED: f32 = 0.45;

static UCT_GREEN: f32 = 0.5;

static UCT_KOMI_MIN_ITERATIONS: usize = 1000;

static TIME_GAP: Time = 100;

pub fn init() {
  unsafe {
    THREADS_COUNT = num_cpus::get();
  }
}

#[inline]
pub fn uct_radius() -> CoordSum {
  UCT_RADIUS
}

#[inline]
pub fn ucb_type() -> UcbType {
  UCB_TYPE
}

#[inline]
pub fn final_ucb_type() -> UcbType {
  FINAL_UCB_TYPE
}

#[inline]
pub fn uct_draw_weight() -> f32 {
  UCT_DRAW_WEIGHT
}

#[inline]
pub fn uctk() -> f32 {
  UCTK
}

#[inline]
pub fn uct_when_create_children() -> usize {
  UCT_WHEN_CREATE_CHILDREN
}

#[inline]
pub fn uct_depth() -> Depth {
  UCT_DEPTH
}

#[inline]
pub fn threads_count() -> usize {
  unsafe { THREADS_COUNT }
}

#[inline]
pub fn uct_komi_type() -> UctKomiType {
  UCT_KOMI_TYPE
}

#[inline]
pub fn uct_red() -> f32 {
  UCT_RED
}

#[inline]
pub fn uct_green() -> f32 {
  UCT_GREEN
}

#[inline]
pub fn uct_komi_min_iterations() -> usize {
  UCT_KOMI_MIN_ITERATIONS
}

#[inline]
pub fn time_gap() -> Time {
  TIME_GAP
}
<commit_msg>Dynamic komi enabled by default.<commit_after>use std::fmt::{Display, Formatter, Result};
use num_cpus;
use types::{CoordSum, Depth, Time};

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum UcbType {
  Winrate,
  Ucb1,
  Ucb1Tuned
}

impl Display for UcbType {
  fn fmt(&self, f: &mut Formatter) -> Result {
    match self {
      &UcbType::Winrate => write!(f, "Winrate"),
      &UcbType::Ucb1 => write!(f, "Ucb1"),
      &UcbType::Ucb1Tuned => write!(f, "Ucb1Tuned")
    }
  }
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum UctKomiType {
  None,
  Static,
  Dynamic
}

impl Display for UctKomiType {
  fn fmt(&self, f: &mut Formatter) -> Result {
    match self {
      &UctKomiType::None => write!(f, "None"),
      &UctKomiType::Static => write!(f, "Static"),
      &UctKomiType::Dynamic => write!(f, "Dynamic")
    }
  }
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum Solver {
  Uct,
  Heuristic
}

impl Display for Solver {
  fn fmt(&self, f: &mut Formatter) -> Result {
    match self {
      &Solver::Uct => write!(f, "Uct"),
      &Solver::Heuristic => write!(f, "Heuristic")
    }
  }
}

static UCT_RADIUS: CoordSum = 3;

static UCB_TYPE: UcbType = UcbType::Ucb1Tuned;

static FINAL_UCB_TYPE: UcbType = UcbType::Winrate;

static UCT_DRAW_WEIGHT: f32 = 0.4;

static UCTK: f32 = 1.0;

static UCT_WHEN_CREATE_CHILDREN: usize = 2;

static UCT_DEPTH: Depth = 8;

static mut THREADS_COUNT: usize = 4;

static UCT_KOMI_TYPE: UctKomiType = UctKomiType::Dynamic;

static UCT_RED: f32 = 0.45;

static UCT_GREEN: f32 = 0.5;

static UCT_KOMI_MIN_ITERATIONS: usize = 1000;

static TIME_GAP: Time = 100;

pub fn init() {
  unsafe {
    THREADS_COUNT = num_cpus::get();
  }
}

#[inline]
pub fn uct_radius() -> CoordSum {
  UCT_RADIUS
}

#[inline]
pub fn ucb_type() -> UcbType {
  UCB_TYPE
}

#[inline]
pub fn final_ucb_type() -> UcbType {
  FINAL_UCB_TYPE
}

#[inline]
pub fn uct_draw_weight() -> f32 {
  UCT_DRAW_WEIGHT
}

#[inline]
pub fn uctk() -> f32 {
  UCTK
}

#[inline]
pub fn uct_when_create_children() -> usize {
  UCT_WHEN_CREATE_CHILDREN
}

#[inline]
pub fn uct_depth() -> Depth {
  UCT_DEPTH
}

#[inline]
pub fn threads_count() -> usize {
  unsafe { THREADS_COUNT }
}

#[inline]
pub fn uct_komi_type() -> UctKomiType {
  UCT_KOMI_TYPE
}

#[inline]
pub fn uct_red() -> f32 {
  UCT_RED
}

#[inline]
pub fn uct_green() -> f32 {
  UCT_GREEN
}

#[inline]
pub fn uct_komi_min_iterations() -> usize {
  UCT_KOMI_MIN_ITERATIONS
}

#[inline]
pub fn time_gap() -> Time {
  TIME_GAP
}
<|endoftext|>


<commit_before>//! # Config module.
//!
//! The Config module provides the beamium configuration.
//! It set defaults and then load config from '/etc', local dir and provided path.

use std::fs::File;
use std::io::Read;
use std::io;
use std::fmt;
use std::string::String;
use std::path::Path;
use std::error;
use std::error::Error;
use yaml_rust::{YamlLoader, ScanError};
use cast;
use std::collections::HashMap;
use regex;
use slog;

pub const REST_TIME: u64 = 10;

#[derive(Debug)]
#[derive(Clone)]
/// Config root.
pub struct Config {
    pub scrapers: Vec<Scraper>,
    pub sinks: Vec<Sink>,
    pub labels: HashMap<String, String>,
    pub parameters: Parameters,
}

#[derive(Debug)]
#[derive(Clone)]
/// Scraper config.
pub struct Scraper {
    pub name: String,
    pub url: String,
    pub period: u64,
    pub format: ScraperFormat,
    pub metrics: Option<regex::RegexSet>,
}

#[derive(Debug)]
#[derive(Clone)]
/// Scraper format.
pub enum ScraperFormat {
    Prometheus,
    Sensision,
}

#[derive(Debug)]
#[derive(Clone)]
/// Sink config.
pub struct Sink {
    pub name: String,
    pub url: String,
    pub token: String,
    pub token_header: String,
    pub selector: Option<regex::Regex>,
    pub ttl: u64,
    pub size: u64,
    pub parallel: u64,
}

#[derive(Debug)]
#[derive(Clone)]
/// Parameters config.
pub struct Parameters {
    pub scan_period: u64,
    pub sink_dir: String,
    pub source_dir: String,
    pub batch_size: u64,
    pub batch_count: u64,
    pub log_file: String,
    pub log_level: slog::Level,
    pub syslog: bool,
    pub timeout: u64,
}

#[derive(Debug)]
/// Config Error.
pub enum ConfigError {
    Io(io::Error),
    Yaml(ScanError),
    Regex(regex::Error),
    Format(Box<Error>),
}

impl From<io::Error> for ConfigError {
    fn from(err: io::Error) -> ConfigError {
        ConfigError::Io(err)
    }
}
impl From<ScanError> for ConfigError {
    fn from(err: ScanError) -> ConfigError {
        ConfigError::Yaml(err)
    }
}
impl From<regex::Error> for ConfigError {
    fn from(err: regex::Error) -> ConfigError {
        ConfigError::Regex(err)
    }
}
impl From<Box<Error>> for ConfigError {
    fn from(err: Box<Error>) -> ConfigError {
        ConfigError::Format(err)
    }
}
impl<'a> From<&'a str> for ConfigError {
    fn from(err: &str) -> ConfigError {
        ConfigError::Format(From::from(err))
    }
}
impl From<String> for ConfigError {
    fn from(err: String) -> ConfigError {
        ConfigError::Format(From::from(err))
    }
}

impl fmt::Display for ConfigError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ConfigError::Io(ref err) => err.fmt(f),
            ConfigError::Yaml(ref err) => err.fmt(f),
            ConfigError::Regex(ref err) => err.fmt(f),
            ConfigError::Format(ref err) => err.fmt(f),
        }
    }
}

impl error::Error for ConfigError {
    fn description(&self) -> &str {
        match *self {
            ConfigError::Io(ref err) => err.description(),
            ConfigError::Yaml(ref err) => err.description(),
            ConfigError::Regex(ref err) => err.description(),
            ConfigError::Format(ref err) => err.description(),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            ConfigError::Io(ref err) => Some(err),
            ConfigError::Yaml(ref err) => Some(err),
            ConfigError::Regex(ref err) => Some(err),
            ConfigError::Format(ref err) => Some(err.as_ref()),
        }
    }
}

/// Load config.
///
/// Setup a defaults config and then load config from '/etc', local dir and provided path.
/// Return Err if provided path is not found or if config is unprocessable.
pub fn load_config(config_path: &str) -> Result<Config, ConfigError> {
    // Defaults
    let mut config = Config {
        scrapers: Vec::new(),
        labels: HashMap::new(),
        sinks: Vec::new(),
        parameters: Parameters {
            scan_period: 1000,
            sink_dir: String::from("sinks"),
            source_dir: String::from("sources"),
            batch_size: 200000,
            batch_count: 250,
            log_file: String::from(env!("CARGO_PKG_NAME")) + ".log",
            log_level: slog::Level::Info,
            syslog: false,
            timeout: 300,
        },
    };

    if config_path.is_empty() {
        // Load from etc
        if Path::new("/etc/beamium/config.yaml").exists() {
            try!(load_path("/etc/beamium/config.yaml", &mut config));
        }

        // Load local
        if Path::new("config.yaml").exists() {
            try!(load_path("config.yaml", &mut config));
        }
    } else {
        // Load from provided path
        try!(load_path(config_path, &mut config));
    }

    Ok(config)
}

/// Extend config from file.
fn load_path<P: AsRef<Path>>(file_path: P, config: &mut Config) -> Result<(), ConfigError> {
    let mut file = try!(File::open(file_path));
    let mut contents = String::new();
    try!(file.read_to_string(&mut contents));
    let docs = try!(YamlLoader::load_from_str(&contents));
    let config_scraper_keys = ["sources", "scrapers"];
    for doc in &docs {
        for config_scraper_key in config_scraper_keys.iter() {
            let key = *config_scraper_key;
            if !doc[key].is_badvalue() {
                if "sources" == key {
                    warn!("'sources' is deprecated and will be removed in further revision. \
                        Please use 'scrapers' instead.",)
                }

                let scrapers = try!(doc[key].as_hash().ok_or(format!("{} should be a map", key)));

                for (k, v) in scrapers {
                    let name = try!(k.as_str().ok_or(format!("{} keys should be a string", key)));
                    let url = try!(v["url"]
                                 .as_str()
                                 .ok_or(format!("{}.{}.url is required and should be a string",
                                                key,
                                                name)));
                    let period = try!(v["period"]
                        .as_i64()
                        .ok_or(format!("{}.{}.period is required and should be a number",
                                       key,
                                       name)));
                    let period = try!(cast::u64(period).map_err(|_| {
                        format!("scrapers.{}.period is invalid", name)
                    }));
                    let format = if v["format"].is_badvalue() {
                        ScraperFormat::Prometheus
                    } else {
                        let f = try!(v["format"]
                                         .as_str()
                                         .ok_or(format!("scrapers.{}.format should be a string",
                                                        name)));

                        if f == "prometheus" {
                            ScraperFormat::Prometheus
                        } else if f == "sensision" {
                            ScraperFormat::Sensision
                        } else {
                            return Err(format!("scrapers.{}.format should be 'prometheus' or \
                                                'sensision'",
                                               name)
                                               .into());
                        }
                    };
                    let metrics = if v["metrics"].is_badvalue() {
                        None
                    } else {
                        let mut metrics = Vec::new();
                        let values =
                            try!(v["metrics"].as_vec().ok_or("metrics should be an array"));
                        for v in values {
                            let value = try!(regex::Regex::new(try!(v.as_str()
                                .ok_or(format!("metrics.{} is invalid", name)))));
                            metrics.push(String::from(r"^(\S*)\s") + value.as_str());
                        }

                        Some(try!(regex::RegexSet::new(&metrics)))
                    };

                    config
                        .scrapers
                        .push(Scraper {
                                  name: String::from(name),
                                  url: String::from(url),
                                  period: period,
                                  format: format,
                                  metrics: metrics,
                              })
                }
            }
        }

        if !doc["sinks"].is_badvalue() {
            let sinks = try!(doc["sinks"].as_hash().ok_or("sinks should be a map"));
            for (k, v) in sinks {
                let name = try!(k.as_str().ok_or("sinks keys should be a string"));
                let url = try!(v["url"]
                                   .as_str()
                                   .ok_or(format!("sinks.{}.url is required and should be a string",
                                                  name)));
                let token = try!(v["token"]
                             .as_str()
                             .ok_or(format!("sinks.{}.token is required and should be a string",
                                            name)));
                let token_header = if v["token-header"].is_badvalue() {
                    "X-Warp10-Token"
                } else {
                    try!(v["token-header"]
                             .as_str()
                             .ok_or(format!("sinks.{}.token-header should be a string", name)))
                };

                let selector = if v["selector"].is_badvalue() {
                    None
                } else {
                    Some(try!(regex::Regex::new(format!("^{}",
                                                        try!(v["selector"]
                                                                 .as_str()
                                                                 .ok_or(format!("sinks.{}.selector \
                                                                            is invalid",
                                                                                name))))
                                                        .as_str())))
                };

                let ttl = if v["ttl"].is_badvalue() {
                    3600
                } else {
                    let ttl = try!(v["ttl"]
                                       .as_i64()
                                       .ok_or(format!("sinks.{}.ttl should be a number",
                                                      name)));
                    try!(cast::u64(ttl).map_err(|_| {
                        format!("sinks.{}.ttl should be a positive number", name)
                    }))
                };

                let size = if v["size"].is_badvalue() {
                    1073741824
                } else {
                    let size = try!(v["size"]
                                        .as_i64()
                                        .ok_or(format!("sinks.{}.size should be a number",
                                                       name)));
                    try!(cast::u64(size).map_err(|_| {
                        format!("sinks.{}.size should be a positive number", name)
                    }))
                };

                let parallel = if v["parallel"].is_badvalue() {
                    1
                } else {
                    let parallel = try!(v["parallel"]
                                            .as_i64()
                                            .ok_or(format!("sinks.{}.parallel should be a number",
                                                           name)));
                    try!(cast::u64(parallel).map_err(|_| {
                        format!("sinks.{}.parallel should be a positive number", name)
                    }))
                };

                config
                    .sinks
                    .push(Sink {
                              name: String::from(name),
                              url: String::from(url),
                              token: String::from(token),
                              token_header: String::from(token_header),
                              selector: selector,
                              ttl: ttl,
                              size: size,
                              parallel: parallel,
                          })
            }
        }

        if !doc["labels"].is_badvalue() {
            let labels = try!(doc["labels"].as_hash().ok_or("labels should be a map"));
            for (k, v) in labels {
                let name = try!(k.as_str().ok_or("labels keys should be a string"));
                let value = try!(v.as_str()
                                     .ok_or(format!("labels.{} value should be a string",
                                                    name)));
                config
                    .labels
                    .insert(String::from(name), String::from(value));
            }
        }

        if !doc["parameters"].is_badvalue() {
            if !doc["parameters"]["source-dir"].is_badvalue() {
                let source_dir = try!(doc["parameters"]["source-dir"]
                    .as_str()
                    .ok_or(format!("parameters.source-dir should be a string")));
                config.parameters.source_dir = String::from(source_dir);
            }

            if !doc["parameters"]["sink-dir"].is_badvalue() {
                let sink_dir = try!(doc["parameters"]["sink-dir"]
                    .as_str()
                    .ok_or(format!("parameters.sink-dir should be a string")));
                config.parameters.sink_dir = String::from(sink_dir);
            }

            if !doc["parameters"]["scan-period"].is_badvalue() {
                let scan_period = try!(doc["parameters"]["scan-period"]
                    .as_i64()
                    .ok_or(format!("parameters.scan-period should be a number")));
                let scan_period = try!(cast::u64(scan_period)
                    .map_err(|_| format!("parameters.scan-period is invalid")));
                config.parameters.scan_period = scan_period;
            }

            if !doc["parameters"]["batch-size"].is_badvalue() {
                let batch_size = try!(doc["parameters"]["batch-size"]
                    .as_i64()
                    .ok_or(format!("parameters.batch-size should be a number")));
                let batch_size = try!(cast::u64(batch_size)
                    .map_err(|_| format!("parameters.batch-size is invalid")));
                config.parameters.batch_size = batch_size;
            }

            if !doc["parameters"]["batch-count"].is_badvalue() {
                let batch_count = try!(doc["parameters"]["batch-count"]
                    .as_i64()
                    .ok_or(format!("parameters.batch-count should be a number")));
                let batch_count = try!(cast::u64(batch_count)
                    .map_err(|_| format!("parameters.batch-count is invalid")));
                config.parameters.batch_count = batch_count;
            }

            if !doc["parameters"]["log-file"].is_badvalue() {
                let log_file = try!(doc["parameters"]["log-file"]
                    .as_str()
                    .ok_or(format!("parameters.log-file should be a string")));
                config.parameters.log_file = String::from(log_file);
            }

            if !doc["parameters"]["log-level"].is_badvalue() {
                let log_level = try!(doc["parameters"]["log-level"]
                    .as_i64()
                    .ok_or(format!("parameters.log-level should be a number")));
                let log_level = try!(cast::u64(log_level)
                    .map_err(|_| format!("parameters.log-level is invalid")));
                let log_level = try!(slog::Level::from_usize(log_level as usize)
                                         .ok_or(format!("parameters.log-level is invalid")));
                config.parameters.log_level = log_level;
            }

            if !doc["parameters"]["syslog"].is_badvalue() {
                let syslog = try!(doc["parameters"]["syslog"]
                                      .as_bool()
                                      .ok_or("parameters.bool should be a boolean"));
                config.parameters.syslog = syslog;

            }
            if !doc["parameters"]["timeout"].is_badvalue() {
                let timeout = try!(doc["parameters"]["timeout"]
                                       .as_i64()
                                       .ok_or("parameters.timeout should be a number"));
                let timeout = try!(cast::u64(timeout)
                                       .map_err(|_| format!("parameters.timeout is invalid")));
                config.parameters.timeout = timeout;
            }
        }
    }
    Ok(())
}
<commit_msg>Fix double load of config file<commit_after>//! # Config module.
//!
//! The Config module provides the beamium configuration.
//! It set defaults and then load config from '/etc', local dir and provided path.

use std::fs::File;
use std::io::Read;
use std::io;
use std::fmt;
use std::string::String;
use std::path::Path;
use std::error;
use std::error::Error;
use yaml_rust::{YamlLoader, ScanError};
use cast;
use std::collections::HashMap;
use regex;
use slog;

pub const REST_TIME: u64 = 10;

#[derive(Debug)]
#[derive(Clone)]
/// Config root.
pub struct Config {
    pub scrapers: Vec<Scraper>,
    pub sinks: Vec<Sink>,
    pub labels: HashMap<String, String>,
    pub parameters: Parameters,
}

#[derive(Debug)]
#[derive(Clone)]
/// Scraper config.
pub struct Scraper {
    pub name: String,
    pub url: String,
    pub period: u64,
    pub format: ScraperFormat,
    pub metrics: Option<regex::RegexSet>,
}

#[derive(Debug)]
#[derive(Clone)]
/// Scraper format.
pub enum ScraperFormat {
    Prometheus,
    Sensision,
}

#[derive(Debug)]
#[derive(Clone)]
/// Sink config.
pub struct Sink {
    pub name: String,
    pub url: String,
    pub token: String,
    pub token_header: String,
    pub selector: Option<regex::Regex>,
    pub ttl: u64,
    pub size: u64,
    pub parallel: u64,
}

#[derive(Debug)]
#[derive(Clone)]
/// Parameters config.
pub struct Parameters {
    pub scan_period: u64,
    pub sink_dir: String,
    pub source_dir: String,
    pub batch_size: u64,
    pub batch_count: u64,
    pub log_file: String,
    pub log_level: slog::Level,
    pub syslog: bool,
    pub timeout: u64,
}

#[derive(Debug)]
/// Config Error.
pub enum ConfigError {
    Io(io::Error),
    Yaml(ScanError),
    Regex(regex::Error),
    Format(Box<Error>),
}

impl From<io::Error> for ConfigError {
    fn from(err: io::Error) -> ConfigError {
        ConfigError::Io(err)
    }
}
impl From<ScanError> for ConfigError {
    fn from(err: ScanError) -> ConfigError {
        ConfigError::Yaml(err)
    }
}
impl From<regex::Error> for ConfigError {
    fn from(err: regex::Error) -> ConfigError {
        ConfigError::Regex(err)
    }
}
impl From<Box<Error>> for ConfigError {
    fn from(err: Box<Error>) -> ConfigError {
        ConfigError::Format(err)
    }
}
impl<'a> From<&'a str> for ConfigError {
    fn from(err: &str) -> ConfigError {
        ConfigError::Format(From::from(err))
    }
}
impl From<String> for ConfigError {
    fn from(err: String) -> ConfigError {
        ConfigError::Format(From::from(err))
    }
}

impl fmt::Display for ConfigError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ConfigError::Io(ref err) => err.fmt(f),
            ConfigError::Yaml(ref err) => err.fmt(f),
            ConfigError::Regex(ref err) => err.fmt(f),
            ConfigError::Format(ref err) => err.fmt(f),
        }
    }
}

impl error::Error for ConfigError {
    fn description(&self) -> &str {
        match *self {
            ConfigError::Io(ref err) => err.description(),
            ConfigError::Yaml(ref err) => err.description(),
            ConfigError::Regex(ref err) => err.description(),
            ConfigError::Format(ref err) => err.description(),
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            ConfigError::Io(ref err) => Some(err),
            ConfigError::Yaml(ref err) => Some(err),
            ConfigError::Regex(ref err) => Some(err),
            ConfigError::Format(ref err) => Some(err.as_ref()),
        }
    }
}

/// Load config.
///
/// Setup a defaults config and then load config from '/etc', local dir and provided path.
/// Return Err if provided path is not found or if config is unprocessable.
pub fn load_config(config_path: &str) -> Result<Config, ConfigError> {
    // Defaults
    let mut config = Config {
        scrapers: Vec::new(),
        labels: HashMap::new(),
        sinks: Vec::new(),
        parameters: Parameters {
            scan_period: 1000,
            sink_dir: String::from("sinks"),
            source_dir: String::from("sources"),
            batch_size: 200000,
            batch_count: 250,
            log_file: String::from(env!("CARGO_PKG_NAME")) + ".log",
            log_level: slog::Level::Info,
            syslog: false,
            timeout: 300,
        },
    };

    if config_path.is_empty() {
        // Load from etc
        if Path::new("/etc/beamium/config.yaml").exists() {
            try!(load_path("/etc/beamium/config.yaml", &mut config));
        } else if Path::new("config.yaml").exists() {
            // Load local
            try!(load_path("config.yaml", &mut config));
        }
    } else {
        // Load from provided path
        try!(load_path(config_path, &mut config));
    }

    Ok(config)
}

/// Extend config from file.
fn load_path<P: AsRef<Path>>(file_path: P, config: &mut Config) -> Result<(), ConfigError> {
    let mut file = try!(File::open(file_path));
    let mut contents = String::new();
    try!(file.read_to_string(&mut contents));
    let docs = try!(YamlLoader::load_from_str(&contents));
    let config_scraper_keys = ["sources", "scrapers"];
    for doc in &docs {
        for config_scraper_key in config_scraper_keys.iter() {
            let key = *config_scraper_key;
            if !doc[key].is_badvalue() {
                if "sources" == key {
                    warn!("'sources' is deprecated and will be removed in further revision. \
                        Please use 'scrapers' instead.",)
                }

                let scrapers = try!(doc[key].as_hash().ok_or(format!("{} should be a map", key)));

                for (k, v) in scrapers {
                    let name = try!(k.as_str().ok_or(format!("{} keys should be a string", key)));
                    let url = try!(v["url"]
                                 .as_str()
                                 .ok_or(format!("{}.{}.url is required and should be a string",
                                                key,
                                                name)));
                    let period = try!(v["period"]
                        .as_i64()
                        .ok_or(format!("{}.{}.period is required and should be a number",
                                       key,
                                       name)));
                    let period = try!(cast::u64(period).map_err(|_| {
                        format!("scrapers.{}.period is invalid", name)
                    }));
                    let format = if v["format"].is_badvalue() {
                        ScraperFormat::Prometheus
                    } else {
                        let f = try!(v["format"]
                                         .as_str()
                                         .ok_or(format!("scrapers.{}.format should be a string",
                                                        name)));

                        if f == "prometheus" {
                            ScraperFormat::Prometheus
                        } else if f == "sensision" {
                            ScraperFormat::Sensision
                        } else {
                            return Err(format!("scrapers.{}.format should be 'prometheus' or \
                                                'sensision'",
                                               name)
                                               .into());
                        }
                    };
                    let metrics = if v["metrics"].is_badvalue() {
                        None
                    } else {
                        let mut metrics = Vec::new();
                        let values =
                            try!(v["metrics"].as_vec().ok_or("metrics should be an array"));
                        for v in values {
                            let value = try!(regex::Regex::new(try!(v.as_str()
                                .ok_or(format!("metrics.{} is invalid", name)))));
                            metrics.push(String::from(r"^(\S*)\s") + value.as_str());
                        }

                        Some(try!(regex::RegexSet::new(&metrics)))
                    };

                    config
                        .scrapers
                        .push(Scraper {
                                  name: String::from(name),
                                  url: String::from(url),
                                  period: period,
                                  format: format,
                                  metrics: metrics,
                              })
                }
            }
        }

        if !doc["sinks"].is_badvalue() {
            let sinks = try!(doc["sinks"].as_hash().ok_or("sinks should be a map"));
            for (k, v) in sinks {
                let name = try!(k.as_str().ok_or("sinks keys should be a string"));
                let url = try!(v["url"]
                                   .as_str()
                                   .ok_or(format!("sinks.{}.url is required and should be a string",
                                                  name)));
                let token = try!(v["token"]
                             .as_str()
                             .ok_or(format!("sinks.{}.token is required and should be a string",
                                            name)));
                let token_header = if v["token-header"].is_badvalue() {
                    "X-Warp10-Token"
                } else {
                    try!(v["token-header"]
                             .as_str()
                             .ok_or(format!("sinks.{}.token-header should be a string", name)))
                };

                let selector = if v["selector"].is_badvalue() {
                    None
                } else {
                    Some(try!(regex::Regex::new(format!("^{}",
                                                        try!(v["selector"]
                                                                 .as_str()
                                                                 .ok_or(format!("sinks.{}.selector \
                                                                            is invalid",
                                                                                name))))
                                                        .as_str())))
                };

                let ttl = if v["ttl"].is_badvalue() {
                    3600
                } else {
                    let ttl = try!(v["ttl"]
                                       .as_i64()
                                       .ok_or(format!("sinks.{}.ttl should be a number",
                                                      name)));
                    try!(cast::u64(ttl).map_err(|_| {
                        format!("sinks.{}.ttl should be a positive number", name)
                    }))
                };

                let size = if v["size"].is_badvalue() {
                    1073741824
                } else {
                    let size = try!(v["size"]
                                        .as_i64()
                                        .ok_or(format!("sinks.{}.size should be a number",
                                                       name)));
                    try!(cast::u64(size).map_err(|_| {
                        format!("sinks.{}.size should be a positive number", name)
                    }))
                };

                let parallel = if v["parallel"].is_badvalue() {
                    1
                } else {
                    let parallel = try!(v["parallel"]
                                            .as_i64()
                                            .ok_or(format!("sinks.{}.parallel should be a number",
                                                           name)));
                    try!(cast::u64(parallel).map_err(|_| {
                        format!("sinks.{}.parallel should be a positive number", name)
                    }))
                };

                config
                    .sinks
                    .push(Sink {
                              name: String::from(name),
                              url: String::from(url),
                              token: String::from(token),
                              token_header: String::from(token_header),
                              selector: selector,
                              ttl: ttl,
                              size: size,
                              parallel: parallel,
                          })
            }
        }

        if !doc["labels"].is_badvalue() {
            let labels = try!(doc["labels"].as_hash().ok_or("labels should be a map"));
            for (k, v) in labels {
                let name = try!(k.as_str().ok_or("labels keys should be a string"));
                let value = try!(v.as_str()
                                     .ok_or(format!("labels.{} value should be a string",
                                                    name)));
                config
                    .labels
                    .insert(String::from(name), String::from(value));
            }
        }

        if !doc["parameters"].is_badvalue() {
            if !doc["parameters"]["source-dir"].is_badvalue() {
                let source_dir = try!(doc["parameters"]["source-dir"]
                    .as_str()
                    .ok_or(format!("parameters.source-dir should be a string")));
                config.parameters.source_dir = String::from(source_dir);
            }

            if !doc["parameters"]["sink-dir"].is_badvalue() {
                let sink_dir = try!(doc["parameters"]["sink-dir"]
                    .as_str()
                    .ok_or(format!("parameters.sink-dir should be a string")));
                config.parameters.sink_dir = String::from(sink_dir);
            }

            if !doc["parameters"]["scan-period"].is_badvalue() {
                let scan_period = try!(doc["parameters"]["scan-period"]
                    .as_i64()
                    .ok_or(format!("parameters.scan-period should be a number")));
                let scan_period = try!(cast::u64(scan_period)
                    .map_err(|_| format!("parameters.scan-period is invalid")));
                config.parameters.scan_period = scan_period;
            }

            if !doc["parameters"]["batch-size"].is_badvalue() {
                let batch_size = try!(doc["parameters"]["batch-size"]
                    .as_i64()
                    .ok_or(format!("parameters.batch-size should be a number")));
                let batch_size = try!(cast::u64(batch_size)
                    .map_err(|_| format!("parameters.batch-size is invalid")));
                config.parameters.batch_size = batch_size;
            }

            if !doc["parameters"]["batch-count"].is_badvalue() {
                let batch_count = try!(doc["parameters"]["batch-count"]
                    .as_i64()
                    .ok_or(format!("parameters.batch-count should be a number")));
                let batch_count = try!(cast::u64(batch_count)
                    .map_err(|_| format!("parameters.batch-count is invalid")));
                config.parameters.batch_count = batch_count;
            }

            if !doc["parameters"]["log-file"].is_badvalue() {
                let log_file = try!(doc["parameters"]["log-file"]
                    .as_str()
                    .ok_or(format!("parameters.log-file should be a string")));
                config.parameters.log_file = String::from(log_file);
            }

            if !doc["parameters"]["log-level"].is_badvalue() {
                let log_level = try!(doc["parameters"]["log-level"]
                    .as_i64()
                    .ok_or(format!("parameters.log-level should be a number")));
                let log_level = try!(cast::u64(log_level)
                    .map_err(|_| format!("parameters.log-level is invalid")));
                let log_level = try!(slog::Level::from_usize(log_level as usize)
                                         .ok_or(format!("parameters.log-level is invalid")));
                config.parameters.log_level = log_level;
            }

            if !doc["parameters"]["syslog"].is_badvalue() {
                let syslog = try!(doc["parameters"]["syslog"]
                                      .as_bool()
                                      .ok_or("parameters.bool should be a boolean"));
                config.parameters.syslog = syslog;

            }
            if !doc["parameters"]["timeout"].is_badvalue() {
                let timeout = try!(doc["parameters"]["timeout"]
                                       .as_i64()
                                       .ok_or("parameters.timeout should be a number"));
                let timeout = try!(cast::u64(timeout)
                                       .map_err(|_| format!("parameters.timeout is invalid")));
                config.parameters.timeout = timeout;
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>use std::io::File;

use collections::treemap::TreeMap;
use serialize::json::ToJson;

use quire::parse;
use J = serialize::json;

use super::yamlutil::{get_string, get_dict, get_list};

pub struct Command {
    pub run: Option<String>,
    pub container: Option<String>,
    pub accepts_arguments: bool,
}

pub struct Variant {
    pub default: Option<String>,
    pub options: Vec<String>,
}

pub struct Container {
    pub default_command: Option<String>,
    pub wrapper_script: Option<String>,
    pub builder: String,
    pub settings: TreeMap<String, String>,
}

pub struct Config {
    pub commands: TreeMap<String, Command>,
    pub variants: TreeMap<String, Variant>,
    pub containers: TreeMap<String, Container>,
}

fn find_config_path(workdir: &Path) -> Option<Path> {
    let mut dir = workdir.clone();
    loop {
        if dir.join("vagga.yaml").exists() {
            return Some(dir);
        }
        if !dir.pop() {
            return None;
        }
    }
}

pub fn find_config(workdir: &Path) -> Result<(Config, Path), String>{
    let cfg_dir = match find_config_path(workdir) {
        Some(path) => path,
        None => return Err(format!(
            "Config not found in path {}", workdir.display())),
    };
    let filename = cfg_dir.join("vagga.yaml");
    let fname = filename.display();
    let data = match File::open(&filename).read_to_str() {
        Ok(data) => data,
        Err(e) => return Err(format!("{}: {}", fname, e)),
    };
    let json = match parse(data.as_slice(), |doc| {
        return doc.to_json();
    }) {
        Ok(json) => json,
        Err(e) => return Err(format!("{}: {}", fname, e)),
    };

    let mut config = Config {
        commands: TreeMap::new(),
        variants: TreeMap::new(),
        containers: TreeMap::new(),
    };

    let root = match json {
        J::Object(val) => val,
        _ => return Err(format!("{}: root node must be mapping", fname)),
    };

    match root.find(&"commands".to_string()) {
        Some(&J::Object(ref commands)) => {
            for (name, jcmd) in commands.iter() {
                let cmd = Command {
                    run: get_string(jcmd, "run"),
                    container: get_string(jcmd, "container"),
                    accepts_arguments: true,  // TODO(tailhook)
                };
                config.commands.insert(name.clone(), cmd);
            }
        }
        Some(_) => return Err(format!(
            "{}: commands key must be mapping", fname)),
        None => {}
    }

    match root.find(&"containers".to_string()) {
        Some(&J::Object(ref containers)) => {
            for (name, jcont) in containers.iter() {
                let cont = Container {
                    default_command: get_string(jcont, "default-command"),
                    wrapper_script: get_string(jcont, "wrapper-script"),
                    builder: get_string(jcont, "builder")
                             .unwrap_or("nix".to_string()),
                    settings: get_dict(jcont, "settings"),
                };
                config.containers.insert(name.clone(), cont);
            }
        }
        Some(_) => return Err(format!(
            "{}: containers key must be mapping", fname)),
        None => {}
    }

    match root.find(&"variants".to_string()) {
        Some(&J::Object(ref variants)) => {
            for (name, jvar) in variants.iter() {
                let var = Variant {
                    default: get_string(jvar, "default"),
                    options: get_list(jvar, "options"),
                };
                config.variants.insert(name.clone(), var);
            }
        }
        Some(_) => return Err(format!(
            "{}: variants key must be mapping", fname)),
        None => {}
    }

    return Ok((config, cfg_dir));
}
<commit_msg>Implement `.vagga/vagga.yaml` as "backup" configuration file<commit_after>use std::io::File;

use collections::treemap::TreeMap;
use serialize::json::ToJson;

use quire::parse;
use J = serialize::json;

use super::yamlutil::{get_string, get_dict, get_list};

pub struct Command {
    pub run: Option<String>,
    pub container: Option<String>,
    pub accepts_arguments: bool,
}

pub struct Variant {
    pub default: Option<String>,
    pub options: Vec<String>,
}

pub struct Container {
    pub default_command: Option<String>,
    pub wrapper_script: Option<String>,
    pub builder: String,
    pub settings: TreeMap<String, String>,
}

pub struct Config {
    pub commands: TreeMap<String, Command>,
    pub variants: TreeMap<String, Variant>,
    pub containers: TreeMap<String, Container>,
}

fn find_config_path(workdir: &Path) -> Option<(Path, Path)> {
    let mut dir = workdir.clone();
    loop {
        let fname = dir.join_many([".vagga", "vagga.yaml"]);
        if fname.exists() {
            return Some((dir, fname));
        }
        let fname = dir.join("vagga.yaml");
        if fname.exists() {
            return Some((dir, fname));
        }
        if !dir.pop() {
            return None;
        }
    }
}

pub fn find_config(workdir: &Path) -> Result<(Config, Path), String>{
    let (cfg_dir, filename) = match find_config_path(workdir) {
        Some(pair) => pair,
        None => return Err(format!(
            "Config not found in path {}", workdir.display())),
    };
    let fname = filename.display();
    let data = match File::open(&filename).read_to_str() {
        Ok(data) => data,
        Err(e) => return Err(format!("{}: {}", fname, e)),
    };
    let json = match parse(data.as_slice(), |doc| {
        return doc.to_json();
    }) {
        Ok(json) => json,
        Err(e) => return Err(format!("{}: {}", fname, e)),
    };

    let mut config = Config {
        commands: TreeMap::new(),
        variants: TreeMap::new(),
        containers: TreeMap::new(),
    };

    let root = match json {
        J::Object(val) => val,
        _ => return Err(format!("{}: root node must be mapping", fname)),
    };

    match root.find(&"commands".to_string()) {
        Some(&J::Object(ref commands)) => {
            for (name, jcmd) in commands.iter() {
                let cmd = Command {
                    run: get_string(jcmd, "run"),
                    container: get_string(jcmd, "container"),
                    accepts_arguments: true,  // TODO(tailhook)
                };
                config.commands.insert(name.clone(), cmd);
            }
        }
        Some(_) => return Err(format!(
            "{}: commands key must be mapping", fname)),
        None => {}
    }

    match root.find(&"containers".to_string()) {
        Some(&J::Object(ref containers)) => {
            for (name, jcont) in containers.iter() {
                let cont = Container {
                    default_command: get_string(jcont, "default-command"),
                    wrapper_script: get_string(jcont, "wrapper-script"),
                    builder: get_string(jcont, "builder")
                             .unwrap_or("nix".to_string()),
                    settings: get_dict(jcont, "settings"),
                };
                config.containers.insert(name.clone(), cont);
            }
        }
        Some(_) => return Err(format!(
            "{}: containers key must be mapping", fname)),
        None => {}
    }

    match root.find(&"variants".to_string()) {
        Some(&J::Object(ref variants)) => {
            for (name, jvar) in variants.iter() {
                let var = Variant {
                    default: get_string(jvar, "default"),
                    options: get_list(jvar, "options"),
                };
                config.variants.insert(name.clone(), var);
            }
        }
        Some(_) => return Err(format!(
            "{}: variants key must be mapping", fname)),
        None => {}
    }

    return Ok((config, cfg_dir));
}
<|endoftext|>


<commit_before>//! Pool configuration.
use std::default::Default;
use std::time::Duration;

/// A builder for `Config`.
///
/// See the documentation of `Config` for more details about the default value
/// and meaning of the configuration parameters.
#[derive(Clone, Debug)]
pub struct Builder {
    c: Config,
}

impl Builder {
    /// Constructs a new `Builder`.
    ///
    /// Parameters are initialized with their default values.
    #[inline]
    pub fn new() -> Builder {
        Builder {
            c: Default::default(),
        }
    }

    /// Sets `pool_size`.
    ///
    /// # Panics
    ///
    /// Panics if `pool_size` is 0.
    #[inline]
    pub fn pool_size(mut self, pool_size: u32) -> Builder {
        assert!(pool_size > 0, "pool_size must be positive");
        self.c.pool_size = pool_size;
        self
    }

    /// Sets `helper_threads`.
    ///
    /// # Panics
    ///
    /// Panics if `helper_threads` is 0.
    #[inline]
    pub fn helper_threads(mut self, helper_threads: u32) -> Builder {
        assert!(helper_threads > 0, "helper_threads must be positive");
        self.c.helper_threads = helper_threads;
        self
    }

    /// Sets `test_on_check_out`.
    #[inline]
    pub fn test_on_check_out(mut self, test_on_check_out: bool) -> Builder {
        self.c.test_on_check_out = test_on_check_out;
        self
    }

    /// Sets `initialization_fail_fast`.
    #[inline]
    pub fn initialization_fail_fast(mut self, initialization_fail_fast: bool) -> Builder {
        self.c.initialization_fail_fast = initialization_fail_fast;
        self
    }

    /// Sets `connection_timeout`.
    ///
    /// # Panics
    ///
    /// Panics if `connection_timeout` is nonpositive.
    #[inline]
    pub fn connection_timeout(mut self, connection_timeout: Duration) -> Builder {
        assert!(connection_timeout > Duration::zero(), "connection_timeout must be positive");
        self.c.connection_timeout = connection_timeout;
        self
    }

    /// Consumes the `Builder`, turning it into a `Config`.
    #[inline]
    pub fn build(self) -> Config {
        self.c
    }
}

/// A struct specifying the runtime configuration of a pool.
///
/// `Config` implements `Default`, which provides a set of reasonable default
/// values. It can be constructed using a `Builder`.
#[derive(Copy, Clone, Debug)]
pub struct Config {
    pool_size: u32,
    helper_threads: u32,
    test_on_check_out: bool,
    initialization_fail_fast: bool,
    connection_timeout: Duration,
}

impl Default for Config {
    #[inline]
    fn default() -> Config {
        Config {
            pool_size: 10,
            helper_threads: 3,
            test_on_check_out: true,
            initialization_fail_fast: true,
            connection_timeout: Duration::seconds(30),
        }
    }
}

impl Config {
    /// Creates a new `Builder`.
    ///
    /// All parameters are initialized to their default values.
    #[inline]
    pub fn builder() -> Builder {
        Builder::new()
    }
    /// The number of connections managed by the pool.
    ///
    /// Defaults to 10.
    #[inline]
    pub fn pool_size(&self) -> u32 {
        self.pool_size
    }

    /// The number of threads that the pool will use for asynchronous
    /// operations such as connection creation and health checks.
    ///
    /// Defaults to 3.
    #[inline]
    pub fn helper_threads(&self) -> u32 {
        self.helper_threads
    }

    /// If true, the health of a connection will be verified via a call to
    /// `ConnectionManager::is_valid` before it is checked out of the pool.
    ///
    /// Defaults to true.
    #[inline]
    pub fn test_on_check_out(&self) -> bool {
        self.test_on_check_out
    }

    /// If true, `Pool::new` will synchronously initialize its connections,
    /// returning an error if they could not be created.
    ///
    /// Defaults to true.
    #[inline]
    pub fn initialization_fail_fast(&self) -> bool {
        self.initialization_fail_fast
    }

    /// Calls to `Pool::get` will wait this long for a connection to become
    /// available before returning an error.
    ///
    /// Defaults to 30 seconds.
    #[inline]
    pub fn connection_timeout(&self) -> Duration {
        self.connection_timeout
    }
}
<commit_msg>Clarify docs a bit<commit_after>//! Pool configuration.
use std::default::Default;
use std::time::Duration;

/// A builder for `Config`.
///
/// See the documentation of `Config` for more details about the default value
/// and meaning of the configuration parameters.
#[derive(Clone, Debug)]
pub struct Builder {
    c: Config,
}

impl Builder {
    /// Constructs a new `Builder`.
    ///
    /// Parameters are initialized with their default values.
    #[inline]
    pub fn new() -> Builder {
        Builder {
            c: Default::default(),
        }
    }

    /// Sets `pool_size`.
    ///
    /// # Panics
    ///
    /// Panics if `pool_size` is 0.
    #[inline]
    pub fn pool_size(mut self, pool_size: u32) -> Builder {
        assert!(pool_size > 0, "pool_size must be positive");
        self.c.pool_size = pool_size;
        self
    }

    /// Sets `helper_threads`.
    ///
    /// # Panics
    ///
    /// Panics if `helper_threads` is 0.
    #[inline]
    pub fn helper_threads(mut self, helper_threads: u32) -> Builder {
        assert!(helper_threads > 0, "helper_threads must be positive");
        self.c.helper_threads = helper_threads;
        self
    }

    /// Sets `test_on_check_out`.
    #[inline]
    pub fn test_on_check_out(mut self, test_on_check_out: bool) -> Builder {
        self.c.test_on_check_out = test_on_check_out;
        self
    }

    /// Sets `initialization_fail_fast`.
    #[inline]
    pub fn initialization_fail_fast(mut self, initialization_fail_fast: bool) -> Builder {
        self.c.initialization_fail_fast = initialization_fail_fast;
        self
    }

    /// Sets `connection_timeout`.
    ///
    /// # Panics
    ///
    /// Panics if `connection_timeout` is nonpositive.
    #[inline]
    pub fn connection_timeout(mut self, connection_timeout: Duration) -> Builder {
        assert!(connection_timeout > Duration::zero(), "connection_timeout must be positive");
        self.c.connection_timeout = connection_timeout;
        self
    }

    /// Consumes the `Builder`, turning it into a `Config`.
    #[inline]
    pub fn build(self) -> Config {
        self.c
    }
}

/// A struct specifying the runtime configuration of a pool.
///
/// `Config` implements `Default`, which provides a set of reasonable default
/// values. It can be constructed using a `Builder`.
#[derive(Copy, Clone, Debug)]
pub struct Config {
    pool_size: u32,
    helper_threads: u32,
    test_on_check_out: bool,
    initialization_fail_fast: bool,
    connection_timeout: Duration,
}

impl Default for Config {
    #[inline]
    fn default() -> Config {
        Config {
            pool_size: 10,
            helper_threads: 3,
            test_on_check_out: true,
            initialization_fail_fast: true,
            connection_timeout: Duration::seconds(30),
        }
    }
}

impl Config {
    /// Creates a new `Builder` which can be used to construct a customized
    /// `Config`.
    ///
    /// All parameters are initialized to their default values.
    #[inline]
    pub fn builder() -> Builder {
        Builder::new()
    }

    /// The number of connections managed by the pool.
    ///
    /// Defaults to 10.
    #[inline]
    pub fn pool_size(&self) -> u32 {
        self.pool_size
    }

    /// The number of threads that the pool will use for asynchronous
    /// operations such as connection creation and health checks.
    ///
    /// Defaults to 3.
    #[inline]
    pub fn helper_threads(&self) -> u32 {
        self.helper_threads
    }

    /// If true, the health of a connection will be verified via a call to
    /// `ConnectionManager::is_valid` before it is checked out of the pool.
    ///
    /// Defaults to true.
    #[inline]
    pub fn test_on_check_out(&self) -> bool {
        self.test_on_check_out
    }

    /// If true, `Pool::new` will synchronously initialize its connections,
    /// returning an error if they could not be created.
    ///
    /// Defaults to true.
    #[inline]
    pub fn initialization_fail_fast(&self) -> bool {
        self.initialization_fail_fast
    }

    /// Calls to `Pool::get` will wait this long for a connection to become
    /// available before returning an error.
    ///
    /// Defaults to 30 seconds.
    #[inline]
    pub fn connection_timeout(&self) -> Duration {
        self.connection_timeout
    }
}
<|endoftext|>


<commit_before>use toml::{self, Parser, Value};

use std::fs::File;
use std::io::prelude::*;
use std::fmt;

/// Contain the configuration for ElasticSearch.
#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct ESConfig {
  pub host: String,
  pub port: u32,
  pub indexes: Vec<String>
}

impl ESConfig {
  pub fn new() -> ESConfig {
    ESConfig {
      host: "localhost".to_owned(),
      port: 9200,
      indexes: vec!["my_index".to_owned()]
    }
  }
}

impl fmt::Display for ESConfig {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "ElasticSearch on {}:{} (indexes: {})",
      self.host, self.port, self.indexes.join(", "))
  }
}

/// Contain instructions about where Searchspot must
/// listen to for new connections.
#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct HTTPConfig {
  pub host: String,
  pub port: u32
}

impl HTTPConfig {
  pub fn new() -> HTTPConfig {
    HTTPConfig {
      host: "127.0.0.1".to_owned(),
      port: 3000,
    }
  }
}

impl fmt::Display for HTTPConfig {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "Listening on http://{}...", self.host)
  }
}

/// Container for ESConfig and HTTPConfig.
#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct Config {
  pub http: HTTPConfig,
  pub es:   ESConfig
}

impl Config {
  /// Return a new `Config` fill with the default values
  pub fn new() -> Config {
    Config {
      es:   ESConfig::new(),
      http: HTTPConfig::new()
    }
  }

  /// Load, parse and return the configuration file
  /// wrapped inside a `Config`.
  pub fn from_file(path: String) -> Config {
    let config_toml = Config::read_file(path).clone();
    Config::parse(config_toml)
  }

  /// Return a `Config` looking for the parameters
  /// inside the ENV variables. `panic!` if there
  /// are some missing.
  pub fn from_env() -> Config {
    let http_config = HTTPConfig {
      host: env!("HTTP_HOST").to_owned(),
      port: env!("HTTP_PORT").parse::<u32>().unwrap()
    };

    let es_config = ESConfig {
      host: env!("ES_HOST").to_owned(),
      port: env!("ES_PORT").parse::<u32>().unwrap(),
      indexes: vec![env!("ES_INDEX").to_owned()]
    };

    Config { http: http_config, es: es_config }
  }

  /// Read a file from the given path and return its content
  pub fn read_file(path: String) -> Option<String> {
    let mut config_toml = String::new();

    let mut file = match File::open(&path) {
      Ok(file) => file,
      Err(_)   => return None
    };

    file.read_to_string(&mut config_toml)
        .unwrap_or_else(|err| panic!("Error while reading config: [{}]", err));

    Some(config_toml)
  }

  /// Parse given TOML configuration file and return it
  /// wrapped inside a `Config`.
  pub fn parse(config_toml: Option<String>) -> Config {
    if config_toml.is_none() {
      println!("{} {}", "Requested configuration file cannot be found.",
                        "The default configuration will be loaded.\n");
      return Config::new();
    }

    let config_toml_ = config_toml.unwrap();
    let mut parser   = Parser::new(&*config_toml_);
    let     toml     = parser.parse();

    match toml {
      Some(config) => {
        let config = Value::Table(config);
        toml::decode(config).unwrap()
      },
      None => {
        println!("{:?}", parser.errors);
        panic!("Error while parsing the configuration file.");
      }
    }
  }
}

#[cfg(test)]
#[allow(non_upper_case_globals)]
mod tests {
  use config::*;

  const sample_config: &'static str = r#"
    [es]
    host    = "123.0.123.0"
    port    = 9000
    indexes = ["save_meguka"]

    [http]
    host = "1.0.0.127"
    port = 3000
  "#;

  #[test]
  fn test_new() {
    // returns a Config fill with the default hardcoded data
    let config = Config::new();
    assert_eq!(config.es.host,   "localhost".to_owned());
    assert_eq!(config.http.host, "127.0.0.1".to_owned());
  }

  #[test]
  fn test_parse() {
    // returns a Config fill with given TOML configuration file
    let config = Config::parse(Some(sample_config.to_owned()));
    assert_eq!(config.es.host,   "123.0.123.0".to_owned());
    assert_eq!(config.http.host, "1.0.0.127".to_owned());
  }
}
<commit_msg>config: use env::var() instead of var!()<commit_after>use toml::{self, Parser, Value};

use std::fs::File;
use std::io::prelude::*;
use std::fmt;
use std::env;

/// Contain the configuration for ElasticSearch.
#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct ESConfig {
  pub host: String,
  pub port: u32,
  pub indexes: Vec<String>
}

impl ESConfig {
  pub fn new() -> ESConfig {
    ESConfig {
      host: "localhost".to_owned(),
      port: 9200,
      indexes: vec!["my_index".to_owned()]
    }
  }
}

impl fmt::Display for ESConfig {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "ElasticSearch on {}:{} (indexes: {})",
      self.host, self.port, self.indexes.join(", "))
  }
}

/// Contain instructions about where Searchspot must
/// listen to for new connections.
#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct HTTPConfig {
  pub host: String,
  pub port: u32
}

impl HTTPConfig {
  pub fn new() -> HTTPConfig {
    HTTPConfig {
      host: "127.0.0.1".to_owned(),
      port: 3000,
    }
  }
}

impl fmt::Display for HTTPConfig {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "Listening on http://{}...", self.host)
  }
}

/// Container for ESConfig and HTTPConfig.
#[derive(RustcEncodable, RustcDecodable, Debug)]
pub struct Config {
  pub http: HTTPConfig,
  pub es:   ESConfig
}

impl Config {
  /// Return a new `Config` fill with the default values
  pub fn new() -> Config {
    Config {
      es:   ESConfig::new(),
      http: HTTPConfig::new()
    }
  }

  /// Load, parse and return the configuration file
  /// wrapped inside a `Config`.
  pub fn from_file(path: String) -> Config {
    let config_toml = Config::read_file(path).clone();
    Config::parse(config_toml)
  }

  /// Return a `Config` looking for the parameters
  /// inside the ENV variables. `panic!` if there
  /// are some missing.
  pub fn from_env() -> Config {
    let http_config = HTTPConfig {
      host: env::var("HTTP_HOST").unwrap().to_owned(),
      port: env::var("HTTP_PORT").unwrap().parse::<u32>().unwrap()
    };

    let es_config = ESConfig {
      host: env::var("ES_HOST").unwrap().to_owned(),
      port: env::var("ES_PORT").unwrap().parse::<u32>().unwrap(),
      indexes: vec![env::var("ES_INDEX").unwrap().to_owned()]
    };

    Config { http: http_config, es: es_config }
  }

  /// Read a file from the given path and return its content
  pub fn read_file(path: String) -> Option<String> {
    let mut config_toml = String::new();

    let mut file = match File::open(&path) {
      Ok(file) => file,
      Err(_)   => return None
    };

    file.read_to_string(&mut config_toml)
        .unwrap_or_else(|err| panic!("Error while reading config: [{}]", err));

    Some(config_toml)
  }

  /// Parse given TOML configuration file and return it
  /// wrapped inside a `Config`.
  pub fn parse(config_toml: Option<String>) -> Config {
    if config_toml.is_none() {
      println!("{} {}", "Requested configuration file cannot be found.",
                        "The default configuration will be loaded.\n");
      return Config::new();
    }

    let config_toml_ = config_toml.unwrap();
    let mut parser   = Parser::new(&*config_toml_);
    let     toml     = parser.parse();

    match toml {
      Some(config) => {
        let config = Value::Table(config);
        toml::decode(config).unwrap()
      },
      None => {
        println!("{:?}", parser.errors);
        panic!("Error while parsing the configuration file.");
      }
    }
  }
}

#[cfg(test)]
#[allow(non_upper_case_globals)]
mod tests {
  use config::*;

  const sample_config: &'static str = r#"
    [es]
    host    = "123.0.123.0"
    port    = 9000
    indexes = ["save_meguka"]

    [http]
    host = "1.0.0.127"
    port = 3000
  "#;

  #[test]
  fn test_new() {
    // returns a Config fill with the default hardcoded data
    let config = Config::new();
    assert_eq!(config.es.host,   "localhost".to_owned());
    assert_eq!(config.http.host, "127.0.0.1".to_owned());
  }

  #[test]
  fn test_parse() {
    // returns a Config fill with given TOML configuration file
    let config = Config::parse(Some(sample_config.to_owned()));
    assert_eq!(config.es.host,   "123.0.123.0".to_owned());
    assert_eq!(config.http.host, "1.0.0.127".to_owned());
  }
}
<|endoftext|>


<commit_before>// Copyright (c) 2020 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use seed::{attrs, iframe, prelude::*};
use serde_with::with_prefix;
use std::collections::HashMap;

pub static IML_METRICS_DASHBOARD_ID: &str = "8Klek6wZz";
pub static IML_METRICS_DASHBOARD_NAME: &str = "iml-metrics";

with_prefix!(grafana_var "var-");

#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct GrafanaChartData<'a> {
    pub org_id: u16,
    pub refresh: &'a str,
    pub panel_id: u16,
    #[serde(flatten, with = "grafana_var")]
    pub vars: HashMap<String, String>,
}

pub(crate) fn create_chart_params<'a>(
    panel_id: u16,
    vars: impl IntoIterator<Item = (impl ToString, impl ToString)>,
) -> GrafanaChartData<'a> {
    let hm: HashMap<String, String> = vars.into_iter().map(|(x, y)| (x.to_string(), y.to_string())).collect();
    GrafanaChartData {
        org_id: 1,
        refresh: "10s",
        panel_id,
        vars: hm,
    }
}

pub fn no_vars() -> Vec<(String, String)> {
    vec![]
}

/// Create an iframe that loads the specified stratagem chart
pub(crate) fn view<'a, T>(
    dashboard_id: &str,
    dashboard_name: &str,
    chart_data: GrafanaChartData<'a>,
    height: &str,
) -> Node<T> {
    iframe![attrs! {
        At::Src => format!("/grafana/d-solo/{}/{}?{}", dashboard_id, dashboard_name, serde_urlencoded::to_string(chart_data).unwrap()),
        At::Width => "100%",
        At::Height => height,
        "frameborder" => 0
    }]
}
<commit_msg>Make grafana graphs "kiosk" mode (#1670)<commit_after>// Copyright (c) 2020 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use seed::{attrs, iframe, prelude::*};
use serde_with::with_prefix;
use std::collections::HashMap;

pub static IML_METRICS_DASHBOARD_ID: &str = "8Klek6wZz";
pub static IML_METRICS_DASHBOARD_NAME: &str = "iml-metrics";

with_prefix!(grafana_var "var-");

#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct GrafanaChartData<'a> {
    pub org_id: u16,
    pub refresh: &'a str,
    pub panel_id: u16,
    #[serde(flatten, with = "grafana_var")]
    pub vars: HashMap<String, String>,
}

pub(crate) fn create_chart_params<'a>(
    panel_id: u16,
    vars: impl IntoIterator<Item = (impl ToString, impl ToString)>,
) -> GrafanaChartData<'a> {
    let hm: HashMap<String, String> = vars.into_iter().map(|(x, y)| (x.to_string(), y.to_string())).collect();
    GrafanaChartData {
        org_id: 1,
        refresh: "10s",
        panel_id,
        vars: hm,
    }
}

pub fn no_vars() -> Vec<(String, String)> {
    vec![]
}

/// Create an iframe that loads the specified stratagem chart
pub(crate) fn view<'a, T>(
    dashboard_id: &str,
    dashboard_name: &str,
    chart_data: GrafanaChartData<'a>,
    height: &str,
) -> Node<T> {
    iframe![attrs! {
        At::Src => format!("/grafana/d-solo/{}/{}?kiosk&{}", dashboard_id, dashboard_name, serde_urlencoded::to_string(chart_data).unwrap()),
        At::Width => "100%",
        At::Height => height,
        "frameborder" => 0
    }]
}
<|endoftext|>


<commit_before>use super::ssz::decode::{
    decode_length,
    Decodable,
};
use super::utils::hash::canonical_hash;
use super::block::{
    MIN_SSZ_BLOCK_LENGTH,
    MAX_SSZ_BLOCK_LENGTH,
};

#[derive(Debug, PartialEq)]
pub enum BlockValidatorError {
    TooShort,
    TooLong,
    BadPowHash,
    SlotTooLow,
    SlotTooHigh,
}

const LENGTH_BYTES: usize = 4;

/// Allows for reading of block values directly from serialized
/// ssz bytes.
#[derive(Debug, PartialEq)]
pub struct SszBlock<'a> {
    ssz: &'a [u8],
    attestation_len: usize,
    pub len: usize,
}

impl<'a> SszBlock<'a> {
    pub fn from_slice(vec: &'a [u8])
        -> Result<Self, BlockValidatorError>
    {
        let ssz = &vec[..];
        let len = vec.len();
        /*
         * Ensure the SSZ is long enough to be a block.
         */
        if len < MIN_SSZ_BLOCK_LENGTH {
            return Err(BlockValidatorError::TooShort);
        }
        /*
         * Ensure the SSZ slice isn't longer than is possible for a block.
         */
        if len > MAX_SSZ_BLOCK_LENGTH {
            return Err(BlockValidatorError::TooLong);
        }
        /*
         * Determine how many bytes are used to store attestation records
         * and ensure that length is enough to store at least one attestation
         * record.
         */
        let attestation_len = decode_length(ssz, 80, LENGTH_BYTES)
            .map_err(|_| BlockValidatorError::TooShort)?;
        if len < (76 + attestation_len + 96) {
            return Err(BlockValidatorError::TooShort);
        }
        Ok(Self{
            ssz,
            attestation_len,
            len,
        })
    }

    pub fn block_hash(&self) -> Vec<u8> {
        canonical_hash(self.ssz)
    }

    pub fn parent_hash(&self) -> &[u8] {
        &self.ssz[5..37]
    }

    pub fn slot_number(&self) -> u64 {
        u64::ssz_decode(&self.ssz, 32).unwrap().0
    }

    pub fn randao_reveal(&self) -> &[u8] {
        &self.ssz[40..72]
    }

    pub fn attestations(&self) -> &[u8] {
        let start = 72 + LENGTH_BYTES;
        &self.ssz[start..(start + self.attestation_len)]
    }

    pub fn pow_chain_ref(&self) -> &[u8] {
        &self.ssz[(self.len - 96)..(self.len - 64)]
    }

    pub fn act_state_root(&self) -> &[u8] {
        &self.ssz[(self.len - 64)..(self.len - 32)]
    }

    pub fn cry_state_root(&self) -> &[u8] {
        &self.ssz[(self.len - 32)..(self.len)]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::block::Block;
    use super::super::attestation_record::AttestationRecord;
    use super::super::ssz::SszStream;

    fn get_block_ssz(b: &Block) -> Vec<u8> {
        let mut ssz_stream = SszStream::new();
        ssz_stream.append(b);
        ssz_stream.drain()
    }

    #[test]
    fn test_ssz_block_zero_attestation_records() {
        let mut b = Block::zero();
        b.attestations = vec![];
        let ssz = get_block_ssz(&b);

        assert_eq!(
            SszBlock::from_slice(&ssz[..]),
            Err(BlockValidatorError::TooShort)
        );
    }

    #[test]
    fn test_ssz_block_single_attestation_record() {
        let mut b = Block::zero();
        b.attestations = vec![AttestationRecord::zero()];
        let ssz = get_block_ssz(&b);

        assert!(SszBlock::from_slice(&ssz[..]).is_ok());
    }
}
<commit_msg>Finish SszBlock, add tests<commit_after>use super::ssz::decode::{
    decode_length,
    Decodable,
};
use super::utils::hash::canonical_hash;
use super::block::{
    MIN_SSZ_BLOCK_LENGTH,
    MAX_SSZ_BLOCK_LENGTH,
};

#[derive(Debug, PartialEq)]
pub enum BlockValidatorError {
    TooShort,
    TooLong,
    BadPowHash,
    SlotTooLow,
    SlotTooHigh,
}

const LENGTH_BYTES: usize = 4;

/// Allows for reading of block values directly from serialized
/// ssz bytes.
#[derive(Debug, PartialEq)]
pub struct SszBlock<'a> {
    ssz: &'a [u8],
    attestation_len: usize,
    pub len: usize,
}

impl<'a> SszBlock<'a> {
    pub fn from_slice(vec: &'a [u8])
        -> Result<Self, BlockValidatorError>
    {
        let ssz = &vec[..];
        let len = vec.len();
        /*
         * Ensure the SSZ is long enough to be a block.
         */
        if len < MIN_SSZ_BLOCK_LENGTH {
            return Err(BlockValidatorError::TooShort);
        }
        /*
         * Ensure the SSZ slice isn't longer than is possible for a block.
         */
        if len > MAX_SSZ_BLOCK_LENGTH {
            return Err(BlockValidatorError::TooLong);
        }
        /*
         * Determine how many bytes are used to store attestation records
         * and ensure that length is enough to store at least one attestation
         * record.
         */
        let attestation_len = decode_length(ssz, 80, LENGTH_BYTES)
            .map_err(|_| BlockValidatorError::TooShort)?;
        if len < (76 + attestation_len + 96) {
            return Err(BlockValidatorError::TooShort);
        }
        Ok(Self{
            ssz,
            attestation_len,
            len,
        })
    }

    pub fn block_hash(&self) -> Vec<u8> {
        canonical_hash(self.ssz)
    }

    pub fn parent_hash(&self) -> &[u8] {
        &self.ssz[4..36]
    }

    pub fn slot_number(&self) -> u64 {
        /*
         * An error should be unreachable from this decode
         * because we checked the length of the array at
         * the initalization of this struct.
         *
         * If you can make this function panic, please report
         * it to paul@sigmaprime.io
         */
        if let Ok((n, _)) = u64::ssz_decode(&self.ssz, 36) {
            n
        } else {
            unreachable!();
        }
    }

    pub fn randao_reveal(&self) -> &[u8] {
        &self.ssz[48..80]
    }

    pub fn attestations(&self) -> &[u8] {
        let start = 80 + LENGTH_BYTES;
        &self.ssz[start..(start + self.attestation_len)]
    }

    pub fn pow_chain_ref(&self) -> &[u8] {
        let start = self.len - (32 + LENGTH_BYTES + 32 + LENGTH_BYTES + 32);
        &self.ssz[start..(start + 32)]
    }

    pub fn act_state_root(&self) -> &[u8] {
        let start = self.len - (32 + LENGTH_BYTES + 32);
        &self.ssz[start..(start + 32)]
    }

    pub fn cry_state_root(&self) -> &[u8] {
        let start = self.len - 32;
        &self.ssz[start..(start + 32)]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::block::Block;
    use super::super::attestation_record::AttestationRecord;
    use super::super::ssz::SszStream;
    use super::super::utils::types::Hash256;

    fn get_block_ssz(b: &Block) -> Vec<u8> {
        let mut ssz_stream = SszStream::new();
        ssz_stream.append(b);
        ssz_stream.drain()
    }

    fn get_attestation_record_ssz(ar: &AttestationRecord) -> Vec<u8> {
        let mut ssz_stream = SszStream::new();
        ssz_stream.append(ar);
        ssz_stream.drain()
    }

    #[test]
    fn test_ssz_block_zero_attestation_records() {
        let mut b = Block::zero();
        b.attestations = vec![];
        let ssz = get_block_ssz(&b);

        assert_eq!(
            SszBlock::from_slice(&ssz[..]),
            Err(BlockValidatorError::TooShort)
        );
    }

    #[test]
    fn test_ssz_block_single_attestation_record_one_byte_short() {
        let mut b = Block::zero();
        b.attestations = vec![AttestationRecord::zero()];
        let ssz = get_block_ssz(&b);

        assert_eq!(
            SszBlock::from_slice(&ssz[0..(ssz.len() - 1)]),
            Err(BlockValidatorError::TooShort)
        );
    }

    #[test]
    fn test_ssz_block_single_attestation_record() {
        let mut b = Block::zero();
        b.attestations = vec![AttestationRecord::zero()];
        let ssz = get_block_ssz(&b);

        assert!(SszBlock::from_slice(&ssz[..]).is_ok());
    }

    #[test]
    fn test_ssz_block_block_hash() {
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());
        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();
        let hash = ssz_block.block_hash();
        // Note: this hash was not generated by some external program,
        // it was simply printed then copied into the code. This test
        // will tell us if the hash changes, not that it matches some
        // canonical reference.
        let expected_hash = [
            28, 184, 51, 12, 226, 15, 73, 50, 66, 19, 168, 149,
            229, 122, 141, 111, 42, 236, 137, 157, 230, 90, 149,
            58, 145, 52, 47, 62, 158, 131, 46, 147
        ];
        assert_eq!(hash, expected_hash);
    }

    #[test]
    fn test_ssz_block_parent_hash() {
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());
        let reference_hash = Hash256::from([42_u8; 32]);
        block.parent_hash = reference_hash.clone();

        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();

        assert_eq!(ssz_block.parent_hash(), &reference_hash.to_vec()[..]);
    }

    #[test]
    fn test_ssz_block_slot_number() {
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());
        block.slot_number = 42;

        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();

        assert_eq!(ssz_block.slot_number(), 42);
    }

    #[test]
    fn test_ssz_block_randao_reveal() {
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());
        let reference_hash = Hash256::from([42_u8; 32]);
        block.randao_reveal = reference_hash.clone();

        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();

        assert_eq!(ssz_block.randao_reveal(), &reference_hash.to_vec()[..]);
    }

    #[test]
    fn test_ssz_block_attestations() {
        /*
         * Single AttestationRecord
         */
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());

        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();
        let ssz_ar = get_attestation_record_ssz(&AttestationRecord::zero());

        assert_eq!(ssz_block.attestations(), &ssz_ar[..]);

        /*
         * Multiple AttestationRecords
         */
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());
        block.attestations.push(AttestationRecord::zero());

        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();
        let mut ssz_ar = get_attestation_record_ssz(&AttestationRecord::zero());
        ssz_ar.append(&mut get_attestation_record_ssz(&AttestationRecord::zero()));

        assert_eq!(ssz_block.attestations(), &ssz_ar[..]);
    }

    #[test]
    fn test_ssz_block_pow_chain_ref() {
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());
        let reference_hash = Hash256::from([42_u8; 32]);
        block.pow_chain_ref = reference_hash.clone();

        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();

        assert_eq!(ssz_block.pow_chain_ref(), &reference_hash.to_vec()[..]);
    }

    #[test]
    fn test_ssz_block_act_state_root() {
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());
        let reference_hash = Hash256::from([42_u8; 32]);
        block.active_state_root = reference_hash.clone();

        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();

        assert_eq!(ssz_block.act_state_root(), &reference_hash.to_vec()[..]);
    }

    #[test]
    fn test_ssz_block_cry_state_root() {
        let mut block = Block::zero();
        block.attestations.push(AttestationRecord::zero());
        let reference_hash = Hash256::from([42_u8; 32]);
        block.crystallized_state_root = reference_hash.clone();

        let serialized = get_block_ssz(&block);
        let ssz_block = SszBlock::from_slice(&serialized).unwrap();

        assert_eq!(ssz_block.cry_state_root(), &reference_hash.to_vec()[..]);
    }
}
<|endoftext|>


<commit_before>use rustc_serialize::hex::{ToHex, FromHex};
use rustc_serialize::base64::{ToBase64, FromBase64, STANDARD};
use std::str;
use std::io::BufReader;
use std::io::BufWriter;
use std::fs::File;
use std::path::Path;
use std::io::prelude::*;
use std::cmp;
use std::fs::OpenOptions;

pub struct CipherText(Vec<u8>);

impl CipherText {
    pub fn new(bytes: &[u8]) -> CipherText {
        CipherText(bytes.to_vec())
    }

    pub fn from_hex(hex: &str) -> Result<CipherText, Error> {
        let cypher_bytes: Vec<u8> = try!(hex.from_hex()
            .map_err(|e| Error::Hex(e.to_string())));
        Ok(CipherText(cypher_bytes))
    }

    pub fn to_hex(&self) -> String {
        let &CipherText(ref vec_bytes) = self;
        vec_bytes.to_hex()
    }

    pub fn from_b64(b64: &str) -> Result<CipherText, Error> {
        let cypher_bytes: Vec<u8> = try!(b64.from_base64()
            .map_err(|e| Error::Base64(e.to_string())));
        Ok(CipherText(cypher_bytes))
    }

    pub fn to_b64(&self) -> String {
        let &CipherText(ref vec_bytes) = self;
        vec_bytes.to_base64(STANDARD)
    }

    pub fn from_file(path: &str) -> Result<CipherText, Error> {
        let mut text = String::new();
        let file = try!(File::open(Path::new(path)).map_err(|e| Error::File(e.to_string())));
        let in_file = BufReader::new(file);
        for line in in_file.lines() {
            text.push_str(&line.unwrap());
        }
        CipherText::from_b64(&text)
    }

    pub fn to_file(&self, path: &str) -> Result<(), Error> {

        let text_b64 = self.to_b64();
        let mut file = OpenOptions::new()
            .write(true)
            .append(true)
            .create(true)
            .open(path)
            .unwrap();

        let mut i = 0;
        let mut j = cmp::min(text_b64.len(), 80);
        while i != j {
            try!(file.write_all(&text_b64[i..j].as_bytes())
                .map_err(|e| Error::File(e.to_string())));

            try!(file.write_all("\n".as_bytes()).map_err(|e| Error::File(e.to_string())));

            i = j;
            j = cmp::min(text_b64.len(), j + 80);
        }
        Ok(())
    }
}

pub struct PlainText(Vec<u8>);

impl PlainText {
    pub fn from_string(string: &str) -> PlainText {
        PlainText(string.as_bytes().to_vec())
    }

    pub fn from_bytes(bytes: &[u8]) -> PlainText {
        PlainText(bytes.to_vec())
    }

    pub fn from_file(path: &str) -> Result<PlainText, Error> {
        let file = try!(File::open(Path::new(path)).map_err(|e| Error::File(e.to_string())));
        let mut in_file = BufReader::new(file);
        let mut text = String::new();
        try!(in_file.read_to_string(&mut text)
            .map_err(|e| Error::File(e.to_string())));
        Ok(PlainText::from_string(&text))

    }

    pub fn to_file(&self, path: &str) -> Result<(), Error> {
        let &PlainText(ref vec_bytes) = self;
        let mut out_file = BufWriter::new(File::create(Path::new(path)).unwrap());
        try!(out_file.write_all(&vec_bytes).map_err(|e| Error::File(e.to_string())));

        Ok(())
    }

    pub fn to_utf8(&self) -> Result<String, Error> {
        let &PlainText(ref vec_bytes) = self;
        let plain = try!(str::from_utf8(&vec_bytes).map_err(|e| Error::UTF8(e.to_string())));
        Ok(plain.to_string())
    }
}

pub fn encrypt(text: &PlainText, key: &str) -> CipherText {
    let &PlainText(ref bytes) = text;
    let xored: Vec<u8> = repeating_xor(bytes, key.as_bytes());
    CipherText::new(&xored)
}

pub fn decrypt(cypher_text: &CipherText, key: &str) -> Result<PlainText, Error> {
    let &CipherText(ref bytes) = cypher_text;
    let xored: Vec<u8> = repeating_xor(bytes, key.as_bytes());
    Ok(PlainText::from_bytes(&xored))
}

pub fn decrypt_file(input_path: &str, output_path: &str, key: &str) -> Result<(), Error> {
    let cipher = try!(CipherText::from_file(input_path));
    let plain = try!(decrypt(&cipher, key));
    try!(plain.to_file(output_path));
    Ok(())
}

pub fn encrypt_file(input_path: &str, output_path: &str, key: &str) -> Result<(), Error> {
    let plain = try!(PlainText::from_file(input_path));
    let cipher = encrypt(&plain, key);
    try!(cipher.to_file(output_path));
    Ok(())
}

#[derive(Debug)]
pub enum Error {
    Hex(String),
    Base64(String),
    UTF8(String),
    File(String),
}

fn repeating_xor(input: &[u8], key: &[u8]) -> Vec<u8> {
    let mut out: Vec<u8> = Vec::with_capacity(input.len());
    unsafe {
        out.set_len(input.len());
    }
    for (i, in_val) in input.iter().enumerate() {
        out[i] = in_val ^ key[i % key.len()];
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encrypt_and_decrypt() {
        check_text("This is the plain text");
        check_text("123, 456");
        check_text("!£$%^&*@~:");
        check_text("日本語");
    }

    #[test]
    fn test_wrong_key() {
        let text = "This is the plain text";
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let decoded_text = decrypt(&encrypt(&plain_text, key), "wrong_key").unwrap();
        assert!(text != decoded_text.to_utf8().unwrap());
    }

    fn check_text(text: &str) {
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let decoded_text = decrypt(&encrypt(&plain_text, key), key).unwrap();
        assert_eq!(text, decoded_text.to_utf8().unwrap());
    }

}
<commit_msg>Added extra tests<commit_after>use rustc_serialize::hex::{ToHex, FromHex};
use rustc_serialize::base64::{ToBase64, FromBase64, STANDARD};
use std::str;
use std::io::BufReader;
use std::io::BufWriter;
use std::fs::File;
use std::path::Path;
use std::io::prelude::*;
use std::cmp;
use std::fs::OpenOptions;
use std::fs;

pub struct CipherText(Vec<u8>);

impl CipherText {
    pub fn new(bytes: &[u8]) -> CipherText {
        CipherText(bytes.to_vec())
    }

    pub fn from_hex(hex: &str) -> Result<CipherText, Error> {
        let cypher_bytes: Vec<u8> = try!(hex.from_hex()
            .map_err(|e| Error::Hex(e.to_string())));
        Ok(CipherText(cypher_bytes))
    }

    pub fn to_hex(&self) -> String {
        let &CipherText(ref vec_bytes) = self;
        vec_bytes.to_hex()
    }

    pub fn from_b64(b64: &str) -> Result<CipherText, Error> {
        let cypher_bytes: Vec<u8> = try!(b64.from_base64()
            .map_err(|e| Error::Base64(e.to_string())));
        Ok(CipherText(cypher_bytes))
    }

    pub fn to_b64(&self) -> String {
        let &CipherText(ref vec_bytes) = self;
        vec_bytes.to_base64(STANDARD)
    }

    pub fn from_file(path: &str) -> Result<CipherText, Error> {
        let mut text = String::new();
        let file = try!(File::open(Path::new(path)).map_err(|e| Error::File(e.to_string())));
        let in_file = BufReader::new(file);
        for line in in_file.lines() {
            text.push_str(&line.unwrap());
        }
        CipherText::from_b64(&text)
    }

    pub fn to_file(&self, path: &str) -> Result<(), Error> {

        let text_b64 = self.to_b64();

        if file_exists(path) {
            try!(fs::remove_file(path).map_err(|e| Error::File(e.to_string())));
        }

        let mut file = try!(OpenOptions::new()
            .write(true)
            .append(true)
            .create(true)
            .open(path)
            .map_err(|e| Error::File(e.to_string())));

        let mut i = 0;
        let mut j = cmp::min(text_b64.len(), 80);
        while i != j {
            try!(file.write_all(&text_b64[i..j].as_bytes())
                .map_err(|e| Error::File(e.to_string())));

            try!(file.write_all("\n".as_bytes()).map_err(|e| Error::File(e.to_string())));

            i = j;
            j = cmp::min(text_b64.len(), j + 80);
        }
        Ok(())
    }
}

fn file_exists(path: &str) -> bool {
    fs::metadata(path).is_ok()
}

pub struct PlainText(Vec<u8>);

impl PlainText {
    pub fn from_string(string: &str) -> PlainText {
        PlainText(string.as_bytes().to_vec())
    }

    pub fn from_bytes(bytes: &[u8]) -> PlainText {
        PlainText(bytes.to_vec())
    }

    pub fn from_file(path: &str) -> Result<PlainText, Error> {
        let file = try!(File::open(Path::new(path)).map_err(|e| Error::File(e.to_string())));
        let mut in_file = BufReader::new(file);
        let mut text = String::new();
        try!(in_file.read_to_string(&mut text)
            .map_err(|e| Error::File(e.to_string())));
        Ok(PlainText::from_string(&text))

    }

    pub fn to_file(&self, path: &str) -> Result<(), Error> {
        let &PlainText(ref vec_bytes) = self;
        let file = try!(OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(path)
            .map_err(|e| Error::File(e.to_string())));

        let mut out_file = BufWriter::new(file);
        try!(out_file.write_all(&vec_bytes).map_err(|e| Error::File(e.to_string())));

        Ok(())
    }

    pub fn to_utf8(&self) -> Result<String, Error> {
        let &PlainText(ref vec_bytes) = self;
        let plain = try!(str::from_utf8(&vec_bytes).map_err(|e| Error::UTF8(e.to_string())));
        Ok(plain.to_string())
    }
}

pub fn encrypt(text: &PlainText, key: &str) -> CipherText {
    let &PlainText(ref bytes) = text;
    let xored: Vec<u8> = repeating_xor(bytes, key.as_bytes());
    CipherText::new(&xored)
}

pub fn decrypt(cypher_text: &CipherText, key: &str) -> Result<PlainText, Error> {
    let &CipherText(ref bytes) = cypher_text;
    let xored: Vec<u8> = repeating_xor(bytes, key.as_bytes());
    Ok(PlainText::from_bytes(&xored))
}

pub fn decrypt_file(input_path: &str, output_path: &str, key: &str) -> Result<(), Error> {
    let cipher = try!(CipherText::from_file(input_path));
    let plain = try!(decrypt(&cipher, key));
    try!(plain.to_file(output_path));
    Ok(())
}

pub fn encrypt_file(input_path: &str, output_path: &str, key: &str) -> Result<(), Error> {
    let plain = try!(PlainText::from_file(input_path));
    let cipher = encrypt(&plain, key);
    try!(cipher.to_file(output_path));
    Ok(())
}

#[derive(Debug)]
pub enum Error {
    Hex(String),
    Base64(String),
    UTF8(String),
    File(String),
}

fn repeating_xor(input: &[u8], key: &[u8]) -> Vec<u8> {
    let mut out: Vec<u8> = Vec::with_capacity(input.len());
    unsafe {
        out.set_len(input.len());
    }
    for (i, in_val) in input.iter().enumerate() {
        out[i] = in_val ^ key[i % key.len()];
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encrypt_and_decrypt() {
        check_text("This is the plain text\n");
        check_text("123, 456");
        check_text("!£$%^&*@~:");
        check_text("日本語");
    }

    #[test]
    fn test_encrypt_and_decrypt_files() {
        check_text_file("This is the plain text\n");
        check_text_file("123, 456");
        check_text_file("!£$%^&*@~:");
        check_text_file("日本語");
        check_text_file("This is the plain text\nThis is the plain text\nThis is the plain text\n");
    }

    #[test]
    fn test_wrong_key() {
        let text = "This is the plain text";
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let decoded_text = decrypt(&encrypt(&plain_text, key), "wrong_key").unwrap();
        assert!(text != decoded_text.to_utf8().unwrap());
    }

    fn check_text(text: &str) {
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let decoded_text = decrypt(&encrypt(&plain_text, key), key).unwrap();
        assert_eq!(text, decoded_text.to_utf8().unwrap());
    }

    fn check_text_file(text: &str) {
        let plain_text = PlainText::from_string(text);
        let key: &str = "toy";
        let plain_path = "./target/test-plain.txt";
        let cipher_path = "./target/test-cipher.txt";
        let decoded_path = "./target/test-decoded.txt";

        plain_text.to_file(plain_path).unwrap();
        encrypt_file(plain_path, cipher_path, key).unwrap();
        decrypt_file(cipher_path, decoded_path, key).unwrap();
        let decoded_text = PlainText::from_file(decoded_path).unwrap();

        assert_eq!(text, decoded_text.to_utf8().unwrap());
    }

}
<|endoftext|>