<commit_before>use std::{
    env, fs,
    io::{Error as IoError, Read, Write},
    path::PathBuf,
    time::Duration,
};

use ansi_term::{Color, Style};
use app_dirs::{get_app_root, AppDataType};
use log::debug;
use serde_derive::{Deserialize, Serialize};

use crate::{
    error::TealdeerError::{self, ConfigError},
    types::PathSource,
};

pub const CONFIG_FILE_NAME: &str = "config.toml";
pub const MAX_CACHE_AGE: Duration = Duration::from_secs(2_592_000); // 30 days
const DEFAULT_UPDATE_INTERVAL_HOURS: u64 = MAX_CACHE_AGE.as_secs() / 3600; // 30 days

fn default_underline() -> bool {
    false
}

fn default_bold() -> bool {
    false
}

fn default_italic() -> bool {
    false
}

#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "lowercase")]
pub enum RawColor {
    Black,
    Red,
    Green,
    Yellow,
    Blue,
    Purple,
    Cyan,
    White,
    Ansi(u8),
    Rgb { r: u8, g: u8, b: u8 },
}

impl From<RawColor> for Color {
    fn from(raw_color: RawColor) -> Self {
        match raw_color {
            RawColor::Black => Self::Black,
            RawColor::Red => Self::Red,
            RawColor::Green => Self::Green,
            RawColor::Yellow => Self::Yellow,
            RawColor::Blue => Self::Blue,
            RawColor::Purple => Self::Purple,
            RawColor::Cyan => Self::Cyan,
            RawColor::White => Self::White,
            RawColor::Ansi(num) => Self::Fixed(num),
            RawColor::Rgb { r, g, b } => Self::RGB(r, g, b),
        }
    }
}

#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawStyle {
    pub foreground: Option<RawColor>,
    pub background: Option<RawColor>,
    #[serde(default = "default_underline")]
    pub underline: bool,
    #[serde(default = "default_bold")]
    pub bold: bool,
    #[serde(default = "default_italic")]
    pub italic: bool,
}

#[allow(clippy::derivable_impls)] // Explicitly control defaults
impl Default for RawStyle {
    fn default() -> Self {
        Self {
            foreground: None,
            background: None,
            underline: false,
            bold: false,
            italic: false,
        }
    }
}

impl From<RawStyle> for Style {
    fn from(raw_style: RawStyle) -> Self {
        let mut style = Self::default();

        if let Some(foreground) = raw_style.foreground {
            style = style.fg(Color::from(foreground));
        }

        if let Some(background) = raw_style.background {
            style = style.on(Color::from(background));
        }

        if raw_style.underline {
            style = style.underline();
        }

        if raw_style.bold {
            style = style.bold();
        }

        if raw_style.italic {
            style = style.italic();
        }

        style
    }
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawStyleConfig {
    #[serde(default)]
    pub description: RawStyle,
    #[serde(default)]
    pub command_name: RawStyle,
    #[serde(default)]
    pub example_text: RawStyle,
    #[serde(default)]
    pub example_code: RawStyle,
    #[serde(default)]
    pub example_variable: RawStyle,
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawDisplayConfig {
    #[serde(default)]
    pub compact: bool,
    #[serde(default)]
    pub use_pager: bool,
}

/// Serde doesn't support default values yet (tracking issue:
/// <https://github.com/serde-rs/serde/issues/368>), so we need to wrap
/// `DEFAULT_UPDATE_INTERVAL_HOURS` in a function to be able to use
/// `#[serde(default = ...)]`
const fn default_auto_update_interval_hours() -> u64 {
    DEFAULT_UPDATE_INTERVAL_HOURS
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawUpdatesConfig {
    #[serde(default)]
    pub auto_update: bool,
    #[serde(default = "default_auto_update_interval_hours")]
    pub auto_update_interval_hours: u64,
}

impl Default for RawUpdatesConfig {
    fn default() -> Self {
        Self {
            auto_update: false,
            auto_update_interval_hours: DEFAULT_UPDATE_INTERVAL_HOURS,
        }
    }
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawDirectoriesConfig {
    #[serde(default)]
    pub custom_pages_dir: Option<PathBuf>,
}

impl Default for RawDirectoriesConfig {
    fn default() -> Self {
        Self {
            custom_pages_dir: get_app_root(AppDataType::UserData, &crate::APP_INFO)
                .map(|path| path.join("pages"))
                .ok(),
        }
    }
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
struct RawConfig {
    style: RawStyleConfig,
    display: RawDisplayConfig,
    updates: RawUpdatesConfig,
    directories: RawDirectoriesConfig,
}

impl RawConfig {
    fn new() -> Self {
        Self::default()
    }
}

impl Default for RawConfig {
    fn default() -> Self {
        let mut raw_config = RawConfig {
            style: RawStyleConfig::default(),
            display: RawDisplayConfig::default(),
            updates: RawUpdatesConfig::default(),
            directories: RawDirectoriesConfig::default(),
        };

        // Set default config
        raw_config.style.example_text.foreground = Some(RawColor::Green);
        raw_config.style.command_name.foreground = Some(RawColor::Cyan);
        raw_config.style.example_code.foreground = Some(RawColor::Cyan);
        raw_config.style.example_variable.foreground = Some(RawColor::Cyan);
        raw_config.style.example_variable.underline = true;

        raw_config
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct StyleConfig {
    pub description: Style,
    pub command_name: Style,
    pub example_text: Style,
    pub example_code: Style,
    pub example_variable: Style,
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct DisplayConfig {
    pub compact: bool,
    pub use_pager: bool,
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct UpdatesConfig {
    pub auto_update: bool,
    pub auto_update_interval: Duration,
}

#[derive(Clone, Debug, PartialEq)]
pub struct DirectoriesConfig {
    pub custom_pages_dir: Option<PathBuf>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Config {
    pub style: StyleConfig,
    pub display: DisplayConfig,
    pub updates: UpdatesConfig,
    pub directories: DirectoriesConfig,
}

impl From<RawConfig> for Config {
    fn from(raw_config: RawConfig) -> Self {
        Self {
            style: StyleConfig {
                command_name: raw_config.style.command_name.into(),
                description: raw_config.style.description.into(),
                example_text: raw_config.style.example_text.into(),
                example_code: raw_config.style.example_code.into(),
                example_variable: raw_config.style.example_variable.into(),
            },
            display: DisplayConfig {
                compact: raw_config.display.compact,
                use_pager: raw_config.display.use_pager,
            },
            updates: UpdatesConfig {
                auto_update: raw_config.updates.auto_update,
                auto_update_interval: Duration::from_secs(
                    raw_config.updates.auto_update_interval_hours * 3600,
                ),
            },
            directories: DirectoriesConfig {
                custom_pages_dir: raw_config.directories.custom_pages_dir,
            },
        }
    }
}

#[allow(clippy::needless_pass_by_value)]
fn map_io_err_to_config_err(e: IoError) -> TealdeerError {
    ConfigError(format!("Io Error: {}", e))
}

impl Config {
    pub fn load(enable_styles: bool) -> Result<Self, TealdeerError> {
        debug!("Loading config");

        // Determine path
        let (config_file_path, _) = get_config_path()
            .map_err(|e| ConfigError(format!("Could not determine config path: {}", e)))?;

        // Load raw config
        let raw_config: RawConfig = if config_file_path.exists() && config_file_path.is_file() {
            let mut config_file =
                fs::File::open(&config_file_path).map_err(map_io_err_to_config_err)?;
            let mut contents = String::new();
            let _ = config_file
                .read_to_string(&mut contents)
                .map_err(map_io_err_to_config_err)?;
            toml::from_str(&contents).map_err(|err| {
                ConfigError(format!(
                    "Failed to parse config file at {:?}:\n{}",
                    config_file_path, err
                ))
            })?
        } else {
            RawConfig::new()
        };

        // Convert to config
        let mut config = Self::from(raw_config);

        // Potentially override styles
        if !enable_styles {
            config.style = StyleConfig {
                command_name: Style::default(),
                description: Style::default(),
                example_text: Style::default(),
                example_code: Style::default(),
                example_variable: Style::default(),
            };
        }

        Ok(config)
    }
}

/// Return the path to the config directory.
///
/// The config dir path can be overridden using the `TEALDEER_CONFIG_DIR` env
/// variable. Otherwise, the user config directory is returned.
///
/// Note that this function does not verify whether the directory at that
/// location exists, or is a directory.
pub fn get_config_dir() -> Result<(PathBuf, PathSource), TealdeerError> {
    // Allow overriding the config directory by setting the
    // $TEALDEER_CONFIG_DIR env variable.
    if let Ok(value) = env::var("TEALDEER_CONFIG_DIR") {
        return Ok((PathBuf::from(value), PathSource::EnvVar));
    };

    // Otherwise, fall back to the user config directory.
    match get_app_root(AppDataType::UserConfig, &crate::APP_INFO) {
        Ok(dirs) => Ok((dirs, PathSource::OsConvention)),
        Err(_) => Err(ConfigError(
            "Could not determine the user config directory.".into(),
        )),
    }
}

/// Return the path to the config file.
///
/// Note that this function does not verify whether the file at that location
/// exists, or is a file.
pub fn get_config_path() -> Result<(PathBuf, PathSource), TealdeerError> {
    let (config_dir, source) = get_config_dir()?;
    let config_file_path = config_dir.join(CONFIG_FILE_NAME);
    Ok((config_file_path, source))
}

/// Create default config file.
pub fn make_default_config() -> Result<PathBuf, TealdeerError> {
    let (config_dir, _) = get_config_dir()?;

    // Ensure that config directory exists
    if !config_dir.exists() {
        if let Err(e) = fs::create_dir_all(&config_dir) {
            return Err(ConfigError(format!(
                "Could not create config directory: {}",
                e
            )));
        }
    } else if !config_dir.is_dir() {
        return Err(ConfigError(format!(
            "Config directory could not be created: {} already exists but is not a directory",
            config_dir.to_string_lossy(),
        )));
    }

    // Ensure that a config file doesn't get overwritten
    let config_file_path = config_dir.join(CONFIG_FILE_NAME);
    if config_file_path.is_file() {
        return Err(ConfigError(format!(
            "A configuration file already exists at {}, no action was taken.",
            config_file_path.to_str().unwrap()
        )));
    }

    // Create default config
    let serialized_config = toml::to_string(&RawConfig::new())
        .map_err(|err| ConfigError(format!("Failed to serialize default config: {}", err)))?;

    // Write default config
    let mut config_file = fs::File::create(&config_file_path).map_err(map_io_err_to_config_err)?;
    let _wc = config_file
        .write(serialized_config.as_bytes())
        .map_err(map_io_err_to_config_err)?;

    Ok(config_file_path)
}

#[test]
fn test_serialize_deserialize() {
    let raw_config = RawConfig::new();
    let serialized = toml::to_string(&raw_config).unwrap();
    let deserialized: RawConfig = toml::from_str(&serialized).unwrap();
    assert_eq!(raw_config, deserialized);
}
<commit_msg>Ensure that default custom_pages_dir ends with trailing slash<commit_after>use std::{
    env, fs,
    io::{Error as IoError, Read, Write},
    path::PathBuf,
    time::Duration,
};

use ansi_term::{Color, Style};
use app_dirs::{get_app_root, AppDataType};
use log::debug;
use serde_derive::{Deserialize, Serialize};

use crate::{
    error::TealdeerError::{self, ConfigError},
    types::PathSource,
};

pub const CONFIG_FILE_NAME: &str = "config.toml";
pub const MAX_CACHE_AGE: Duration = Duration::from_secs(2_592_000); // 30 days
const DEFAULT_UPDATE_INTERVAL_HOURS: u64 = MAX_CACHE_AGE.as_secs() / 3600; // 30 days

fn default_underline() -> bool {
    false
}

fn default_bold() -> bool {
    false
}

fn default_italic() -> bool {
    false
}

#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "lowercase")]
pub enum RawColor {
    Black,
    Red,
    Green,
    Yellow,
    Blue,
    Purple,
    Cyan,
    White,
    Ansi(u8),
    Rgb { r: u8, g: u8, b: u8 },
}

impl From<RawColor> for Color {
    fn from(raw_color: RawColor) -> Self {
        match raw_color {
            RawColor::Black => Self::Black,
            RawColor::Red => Self::Red,
            RawColor::Green => Self::Green,
            RawColor::Yellow => Self::Yellow,
            RawColor::Blue => Self::Blue,
            RawColor::Purple => Self::Purple,
            RawColor::Cyan => Self::Cyan,
            RawColor::White => Self::White,
            RawColor::Ansi(num) => Self::Fixed(num),
            RawColor::Rgb { r, g, b } => Self::RGB(r, g, b),
        }
    }
}

#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawStyle {
    pub foreground: Option<RawColor>,
    pub background: Option<RawColor>,
    #[serde(default = "default_underline")]
    pub underline: bool,
    #[serde(default = "default_bold")]
    pub bold: bool,
    #[serde(default = "default_italic")]
    pub italic: bool,
}

#[allow(clippy::derivable_impls)] // Explicitly control defaults
impl Default for RawStyle {
    fn default() -> Self {
        Self {
            foreground: None,
            background: None,
            underline: false,
            bold: false,
            italic: false,
        }
    }
}

impl From<RawStyle> for Style {
    fn from(raw_style: RawStyle) -> Self {
        let mut style = Self::default();

        if let Some(foreground) = raw_style.foreground {
            style = style.fg(Color::from(foreground));
        }

        if let Some(background) = raw_style.background {
            style = style.on(Color::from(background));
        }

        if raw_style.underline {
            style = style.underline();
        }

        if raw_style.bold {
            style = style.bold();
        }

        if raw_style.italic {
            style = style.italic();
        }

        style
    }
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawStyleConfig {
    #[serde(default)]
    pub description: RawStyle,
    #[serde(default)]
    pub command_name: RawStyle,
    #[serde(default)]
    pub example_text: RawStyle,
    #[serde(default)]
    pub example_code: RawStyle,
    #[serde(default)]
    pub example_variable: RawStyle,
}

#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
struct RawDisplayConfig {
    #[serde(default)]
    pub compact: bool,
    #[serde(default)]
    pub use_pager: bool,
}

/// Serde doesn't support default values yet (tracking issue:
/// <https://github.com/serde-rs/serde/issues/368>), so we need to wrap
/// `DEFAULT_UPDATE_INTERVAL_HOURS` in a function to be able to use
/// `#[serde(default = ...)]`
const fn default_auto_update_interval_hours() -> u64 {
    DEFAULT_UPDATE_INTERVAL_HOURS
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawUpdatesConfig {
    #[serde(default)]
    pub auto_update: bool,
    #[serde(default = "default_auto_update_interval_hours")]
    pub auto_update_interval_hours: u64,
}

impl Default for RawUpdatesConfig {
    fn default() -> Self {
        Self {
            auto_update: false,
            auto_update_interval_hours: DEFAULT_UPDATE_INTERVAL_HOURS,
        }
    }
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct RawDirectoriesConfig {
    #[serde(default)]
    pub custom_pages_dir: Option<PathBuf>,
}

impl Default for RawDirectoriesConfig {
    fn default() -> Self {
        Self {
            custom_pages_dir: get_app_root(AppDataType::UserData, &crate::APP_INFO)
                .map(|path| {
                    // Note: The `join("")` call ensures that there's a trailing slash
                    path.join("pages").join("")
                })
                .ok(),
        }
    }
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
struct RawConfig {
    style: RawStyleConfig,
    display: RawDisplayConfig,
    updates: RawUpdatesConfig,
    directories: RawDirectoriesConfig,
}

impl RawConfig {
    fn new() -> Self {
        Self::default()
    }
}

impl Default for RawConfig {
    fn default() -> Self {
        let mut raw_config = RawConfig {
            style: RawStyleConfig::default(),
            display: RawDisplayConfig::default(),
            updates: RawUpdatesConfig::default(),
            directories: RawDirectoriesConfig::default(),
        };

        // Set default config
        raw_config.style.example_text.foreground = Some(RawColor::Green);
        raw_config.style.command_name.foreground = Some(RawColor::Cyan);
        raw_config.style.example_code.foreground = Some(RawColor::Cyan);
        raw_config.style.example_variable.foreground = Some(RawColor::Cyan);
        raw_config.style.example_variable.underline = true;

        raw_config
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct StyleConfig {
    pub description: Style,
    pub command_name: Style,
    pub example_text: Style,
    pub example_code: Style,
    pub example_variable: Style,
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct DisplayConfig {
    pub compact: bool,
    pub use_pager: bool,
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct UpdatesConfig {
    pub auto_update: bool,
    pub auto_update_interval: Duration,
}

#[derive(Clone, Debug, PartialEq)]
pub struct DirectoriesConfig {
    pub custom_pages_dir: Option<PathBuf>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Config {
    pub style: StyleConfig,
    pub display: DisplayConfig,
    pub updates: UpdatesConfig,
    pub directories: DirectoriesConfig,
}

impl From<RawConfig> for Config {
    fn from(raw_config: RawConfig) -> Self {
        Self {
            style: StyleConfig {
                command_name: raw_config.style.command_name.into(),
                description: raw_config.style.description.into(),
                example_text: raw_config.style.example_text.into(),
                example_code: raw_config.style.example_code.into(),
                example_variable: raw_config.style.example_variable.into(),
            },
            display: DisplayConfig {
                compact: raw_config.display.compact,
                use_pager: raw_config.display.use_pager,
            },
            updates: UpdatesConfig {
                auto_update: raw_config.updates.auto_update,
                auto_update_interval: Duration::from_secs(
                    raw_config.updates.auto_update_interval_hours * 3600,
                ),
            },
            directories: DirectoriesConfig {
                custom_pages_dir: raw_config.directories.custom_pages_dir,
            },
        }
    }
}

#[allow(clippy::needless_pass_by_value)]
fn map_io_err_to_config_err(e: IoError) -> TealdeerError {
    ConfigError(format!("Io Error: {}", e))
}

impl Config {
    pub fn load(enable_styles: bool) -> Result<Self, TealdeerError> {
        debug!("Loading config");

        // Determine path
        let (config_file_path, _) = get_config_path()
            .map_err(|e| ConfigError(format!("Could not determine config path: {}", e)))?;

        // Load raw config
        let raw_config: RawConfig = if config_file_path.exists() && config_file_path.is_file() {
            let mut config_file =
                fs::File::open(&config_file_path).map_err(map_io_err_to_config_err)?;
            let mut contents = String::new();
            let _ = config_file
                .read_to_string(&mut contents)
                .map_err(map_io_err_to_config_err)?;
            toml::from_str(&contents).map_err(|err| {
                ConfigError(format!(
                    "Failed to parse config file at {:?}:\n{}",
                    config_file_path, err
                ))
            })?
        } else {
            RawConfig::new()
        };

        // Convert to config
        let mut config = Self::from(raw_config);

        // Potentially override styles
        if !enable_styles {
            config.style = StyleConfig {
                command_name: Style::default(),
                description: Style::default(),
                example_text: Style::default(),
                example_code: Style::default(),
                example_variable: Style::default(),
            };
        }

        Ok(config)
    }
}

/// Return the path to the config directory.
///
/// The config dir path can be overridden using the `TEALDEER_CONFIG_DIR` env
/// variable. Otherwise, the user config directory is returned.
///
/// Note that this function does not verify whether the directory at that
/// location exists, or is a directory.
pub fn get_config_dir() -> Result<(PathBuf, PathSource), TealdeerError> {
    // Allow overriding the config directory by setting the
    // $TEALDEER_CONFIG_DIR env variable.
    if let Ok(value) = env::var("TEALDEER_CONFIG_DIR") {
        return Ok((PathBuf::from(value), PathSource::EnvVar));
    };

    // Otherwise, fall back to the user config directory.
    match get_app_root(AppDataType::UserConfig, &crate::APP_INFO) {
        Ok(dirs) => Ok((dirs, PathSource::OsConvention)),
        Err(_) => Err(ConfigError(
            "Could not determine the user config directory.".into(),
        )),
    }
}

/// Return the path to the config file.
///
/// Note that this function does not verify whether the file at that location
/// exists, or is a file.
pub fn get_config_path() -> Result<(PathBuf, PathSource), TealdeerError> {
    let (config_dir, source) = get_config_dir()?;
    let config_file_path = config_dir.join(CONFIG_FILE_NAME);
    Ok((config_file_path, source))
}

/// Create default config file.
pub fn make_default_config() -> Result<PathBuf, TealdeerError> {
    let (config_dir, _) = get_config_dir()?;

    // Ensure that config directory exists
    if !config_dir.exists() {
        if let Err(e) = fs::create_dir_all(&config_dir) {
            return Err(ConfigError(format!(
                "Could not create config directory: {}",
                e
            )));
        }
    } else if !config_dir.is_dir() {
        return Err(ConfigError(format!(
            "Config directory could not be created: {} already exists but is not a directory",
            config_dir.to_string_lossy(),
        )));
    }

    // Ensure that a config file doesn't get overwritten
    let config_file_path = config_dir.join(CONFIG_FILE_NAME);
    if config_file_path.is_file() {
        return Err(ConfigError(format!(
            "A configuration file already exists at {}, no action was taken.",
            config_file_path.to_str().unwrap()
        )));
    }

    // Create default config
    let serialized_config = toml::to_string(&RawConfig::new())
        .map_err(|err| ConfigError(format!("Failed to serialize default config: {}", err)))?;

    // Write default config
    let mut config_file = fs::File::create(&config_file_path).map_err(map_io_err_to_config_err)?;
    let _wc = config_file
        .write(serialized_config.as_bytes())
        .map_err(map_io_err_to_config_err)?;

    Ok(config_file_path)
}

#[test]
fn test_serialize_deserialize() {
    let raw_config = RawConfig::new();
    let serialized = toml::to_string(&raw_config).unwrap();
    let deserialized: RawConfig = toml::from_str(&serialized).unwrap();
    assert_eq!(raw_config, deserialized);
}
<|endoftext|>


<commit_before>extern crate rustc_serialize;

use std::fs::File;
use std::io::prelude::*;
use toml::{Parser, Value};
use toml;

#[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!["honeypot_dev_talents".to_owned()]
    }
  }
}

#[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,
    }
  }
}

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

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 load_config(path: String) -> Config {
    let config_toml = Config::read_file(path).clone();
    Config::parse(config_toml)
  }

  /// 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 = ["honeypot_dev_talents"]

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

  #[test]
  fn test_config() {
    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() {
    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: document the structs<commit_after>extern crate rustc_serialize;

use std::fs::File;
use std::io::prelude::*;
use toml::{Parser, Value};
use toml;

/// 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!["honeypot_dev_talents".to_owned()]
    }
  }
}

/// Contain instructions about where Honeysearch 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,
    }
  }
}

/// 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 load_config(path: String) -> Config {
    let config_toml = Config::read_file(path).clone();
    Config::parse(config_toml)
  }

  /// 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 = ["honeypot_dev_talents"]

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

  #[test]
  fn test_config() {
    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() {
    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>use std::path::{Path, PathBuf};
use std::fs::File;
use std::io::Read;
use toml::{self, Value};
use toml::value::Table;
use serde::{Deserialize, Deserializer};

use errors::*;

/// The overall configuration object for MDBook.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct Config {
    /// Metadata about the book.
    pub book: BookConfig,
    rest: Table,
}

impl Config {
    /// Load a `Config` from some string.
    pub fn from_str(src: &str) -> Result<Config> {
        toml::from_str(src).chain_err(|| Error::from("Invalid configuration file"))
    }

    /// Load the configuration file from disk.
    pub fn from_disk<P: AsRef<Path>>(config_file: P) -> Result<Config> {
        let mut buffer = String::new();
        File::open(config_file).chain_err(|| "Unable to open the configuration file")?
                               .read_to_string(&mut buffer)
                               .chain_err(|| "Couldn't read the file")?;

        Config::from_str(&buffer)
    }

    /// Fetch an arbitrary item from the `Config` as a `toml::Value`.
    ///
    /// You can use dotted indices to access nested items (e.g.
    /// `output.html.playpen` will fetch the "playpen" out of the html output
    /// table).
    pub fn get(&self, key: &str) -> Option<&Value> {
        let pieces: Vec<_> = key.split(".").collect();
        recursive_get(&pieces, &self.rest)
    }

    /// Fetch a value from the `Config` so you can mutate it.
    pub fn get_mut<'a>(&'a mut self, key: &str) -> Option<&'a mut Value> {
        let pieces: Vec<_> = key.split(".").collect();
        recursive_get_mut(&pieces, &mut self.rest)
    }

    /// Convenience method for getting the html renderer's configuration.
    ///
    /// # Note
    ///
    /// This is for compatibility only. It will be removed completely once the
    /// rendering and plugin system is established.
    pub fn html_config(&self) -> Option<HtmlConfig> {
        self.get_deserialized("output.html").ok()
    }

    /// Convenience function to fetch a value from the config and deserialize it
    /// into some arbitrary type.
    pub fn get_deserialized<'de, T: Deserialize<'de>, S: AsRef<str>>(&self, name: S) -> Result<T> {
        let name = name.as_ref();

        if let Some(value) = self.get(name) {
            value.clone()
                 .try_into()
                 .chain_err(|| "Couldn't deserialize the value")
        } else {
            bail!("Key not found, {:?}", name)
        }
    }
}

fn recursive_get<'a>(key: &[&str], table: &'a Table) -> Option<&'a Value> {
    if key.is_empty() {
        return None;
    } else if key.len() == 1 {
        return table.get(key[0]);
    }

    let first = key[0];
    let rest = &key[1..];

    if let Some(&Value::Table(ref nested)) = table.get(first) {
        recursive_get(rest, nested)
    } else {
        None
    }
}

fn recursive_get_mut<'a>(key: &[&str], table: &'a mut Table) -> Option<&'a mut Value> {
    // TODO: Figure out how to abstract over mutability to reduce copy-pasta
    if key.is_empty() {
        return None;
    } else if key.len() == 1 {
        return table.get_mut(key[0]);
    }

    let first = key[0];
    let rest = &key[1..];

    if let Some(&mut Value::Table(ref mut nested)) = table.get_mut(first) {
        recursive_get_mut(rest, nested)
    } else {
        None
    }
}

impl<'de> Deserialize<'de> for Config {
    fn deserialize<D: Deserializer<'de>>(de: D) -> ::std::result::Result<Self, D::Error> {
        let raw = Value::deserialize(de)?;
        if let Value::Table(mut table) = raw {
            let book: BookConfig = table.remove("book")
                                        .and_then(|value| value.try_into().ok())
                                        .unwrap_or_default();
            Ok(Config {
                   book: book,
                   rest: table,
               })
        } else {
            use serde::de::Error;
            Err(D::Error::custom("A config file should always be a toml table"))
        }
    }
}

/// Configuration options which are specific to the book and required for
/// loading it from disk.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(default, rename_all = "kebab-case")]
pub struct BookConfig {
    /// The book's title.
    pub title: Option<String>,
    /// The book's authors.
    pub authors: Vec<String>,
    /// An optional description for the book.
    pub description: Option<String>,
    /// Location of the book source relative to the book's root directory.
    pub src: PathBuf,
    /// Where to put built artefacts relative to the book's root directory.
    pub build_dir: PathBuf,
    /// Does this book support more than one language?
    pub multilingual: bool,
}

impl Default for BookConfig {
    fn default() -> BookConfig {
        BookConfig {
            title: None,
            authors: Vec::new(),
            description: None,
            src: PathBuf::from("src"),
            build_dir: PathBuf::from("book"),
            multilingual: false,
        }
    }
}

#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(default, rename_all = "kebab-case")]
pub struct HtmlConfig {
    pub theme: Option<PathBuf>,
    pub curly_quotes: bool,
    pub mathjax_support: bool,
    pub google_analytics: Option<String>,
    pub additional_css: Vec<PathBuf>,
    pub additional_js: Vec<PathBuf>,
    pub playpen: Playpen,
}

/// Configuration for tweaking how the the HTML renderer handles the playpen.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct Playpen {
    pub editor: PathBuf,
    pub editable: bool,
}

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

    const COMPLEX_CONFIG: &'static str = r#"
        [book]
        title = "Some Book"
        authors = ["Michael-F-Bryan <michaelfbryan@gmail.com>"]
        description = "A completely useless book"
        multilingual = true
        src = "source"
        build-dir = "outputs"

        [output.html]
        theme = "./themedir"
        curly-quotes = true
        google-analytics = "123456"
        additional-css = ["./foo/bar/baz.css"]

        [output.html.playpen]
        editable = true
        editor = "ace"
        "#;

    #[test]
    fn load_a_complex_config_file() {
        let src = COMPLEX_CONFIG;

        let book_should_be = BookConfig {
            title: Some(String::from("Some Book")),
            authors: vec![String::from("Michael-F-Bryan <michaelfbryan@gmail.com>")],
            description: Some(String::from("A completely useless book")),
            multilingual: true,
            src: PathBuf::from("source"),
            build_dir: PathBuf::from("outputs"),
            ..Default::default()
        };
        let playpen_should_be = Playpen {
            editable: true,
            editor: PathBuf::from("ace"),
        };
        let html_should_be = HtmlConfig {
            curly_quotes: true,
            google_analytics: Some(String::from("123456")),
            additional_css: vec![PathBuf::from("./foo/bar/baz.css")],
            theme: Some(PathBuf::from("./themedir")),
            playpen: playpen_should_be,
            ..Default::default()
        };

        let got = Config::from_str(src).unwrap();

        assert_eq!(got.book, book_should_be);
        assert_eq!(got.html_config().unwrap(), html_should_be);
    }

    #[test]
    fn load_arbitrary_output_type() {
        #[derive(Debug, Deserialize, PartialEq)]
        struct RandomOutput {
            foo: u32,
            bar: String,
            baz: Vec<bool>,
        }

        let src = r#"
        [output.random]
        foo = 5
        bar = "Hello World"
        baz = [true, true, false]
        "#;

        let should_be = RandomOutput {
            foo: 5,
            bar: String::from("Hello World"),
            baz: vec![true, true, false],
        };

        let cfg = Config::from_str(src).unwrap();
        let got: RandomOutput = cfg.get_deserialized("output.random").unwrap();

        assert_eq!(got, should_be);

        let baz: Vec<bool> = cfg.get_deserialized("output.random.baz").unwrap();
        let baz_should_be = vec![true, true, false];

        assert_eq!(baz, baz_should_be);
    }

#[test]
fn mutate_some_stuff() {
    // really this is just a sanity check to make sure the borrow checker
    // is happy...
    let src = COMPLEX_CONFIG;
    let mut config = Config::from_str(src).unwrap();
    let key = "output.html.playpen.editable";

    assert_eq!(config.get(key).unwrap(), &Value::Boolean(true));
    *config.get_mut(key).unwrap() = Value::Boolean(false);
    assert_eq!(config.get(key).unwrap(), &Value::Boolean(false));
}
}
<commit_msg>Added in legacy config format support<commit_after>use std::path::{Path, PathBuf};
use std::fs::File;
use std::io::Read;
use toml::{self, Value};
use toml::value::Table;
use serde::{Deserialize, Deserializer};

use errors::*;

/// The overall configuration object for MDBook.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct Config {
    /// Metadata about the book.
    pub book: BookConfig,
    rest: Table,
}

impl Config {
    /// Load a `Config` from some string.
    pub fn from_str(src: &str) -> Result<Config> {
        toml::from_str(src).chain_err(|| Error::from("Invalid configuration file"))
    }

    /// Load the configuration file from disk.
    pub fn from_disk<P: AsRef<Path>>(config_file: P) -> Result<Config> {
        let mut buffer = String::new();
        File::open(config_file).chain_err(|| "Unable to open the configuration file")?
                               .read_to_string(&mut buffer)
                               .chain_err(|| "Couldn't read the file")?;

        Config::from_str(&buffer)
    }

    /// Fetch an arbitrary item from the `Config` as a `toml::Value`.
    ///
    /// You can use dotted indices to access nested items (e.g.
    /// `output.html.playpen` will fetch the "playpen" out of the html output
    /// table).
    pub fn get(&self, key: &str) -> Option<&Value> {
        let pieces: Vec<_> = key.split(".").collect();
        recursive_get(&pieces, &self.rest)
    }

    /// Fetch a value from the `Config` so you can mutate it.
    pub fn get_mut<'a>(&'a mut self, key: &str) -> Option<&'a mut Value> {
        let pieces: Vec<_> = key.split(".").collect();
        recursive_get_mut(&pieces, &mut self.rest)
    }

    /// Convenience method for getting the html renderer's configuration.
    ///
    /// # Note
    ///
    /// This is for compatibility only. It will be removed completely once the
    /// rendering and plugin system is established.
    pub fn html_config(&self) -> Option<HtmlConfig> {
        self.get_deserialized("output.html").ok()
    }

    /// Convenience function to fetch a value from the config and deserialize it
    /// into some arbitrary type.
    pub fn get_deserialized<'de, T: Deserialize<'de>, S: AsRef<str>>(&self, name: S) -> Result<T> {
        let name = name.as_ref();

        if let Some(value) = self.get(name) {
            value.clone()
                 .try_into()
                 .chain_err(|| "Couldn't deserialize the value")
        } else {
            bail!("Key not found, {:?}", name)
        }
    }

    fn from_legacy(mut table: Table) -> Config {
        let mut cfg = Config::default();

        // we use a macro here instead of a normal loop because the $out
        // variable can be different types. This way we can make type inference
        // figure out what try_into() deserializes to.
        macro_rules! get_and_insert {
            ($table:expr, $key:expr => $out:expr) => {
                if let Some(value) = $table.remove($key).and_then(|v| v.try_into().ok()) {
                    $out = value;
                }
            };
        }

        get_and_insert!(table, "title" => cfg.book.title);
        get_and_insert!(table, "authors" => cfg.book.authors);
        get_and_insert!(table, "source" => cfg.book.src);
        get_and_insert!(table, "description" => cfg.book.description);

        // This complicated chain of and_then's is so we can move
        // "output.html.destination" to "book.build_dir" and parse it into a
        // PathBuf.
        let destination: Option<PathBuf> = table.get_mut("output")
            .and_then(|output| output.as_table_mut())
            .and_then(|output| output.get_mut("html"))
            .and_then(|html| html.as_table_mut())
            .and_then(|html| html.remove("destination"))
            .and_then(|dest| dest.try_into().ok());

        if let Some(dest) = destination {
            cfg.book.build_dir = dest;
        }

        cfg.rest = table;
        cfg
    }
}

fn recursive_get<'a>(key: &[&str], table: &'a Table) -> Option<&'a Value> {
    if key.is_empty() {
        return None;
    } else if key.len() == 1 {
        return table.get(key[0]);
    }

    let first = key[0];
    let rest = &key[1..];

    if let Some(&Value::Table(ref nested)) = table.get(first) {
        recursive_get(rest, nested)
    } else {
        None
    }
}

fn recursive_get_mut<'a>(key: &[&str], table: &'a mut Table) -> Option<&'a mut Value> {
    // TODO: Figure out how to abstract over mutability to reduce copy-pasta
    if key.is_empty() {
        return None;
    } else if key.len() == 1 {
        return table.get_mut(key[0]);
    }

    let first = key[0];
    let rest = &key[1..];

    if let Some(&mut Value::Table(ref mut nested)) = table.get_mut(first) {
        recursive_get_mut(rest, nested)
    } else {
        None
    }
}

impl<'de> Deserialize<'de> for Config {
    fn deserialize<D: Deserializer<'de>>(de: D) -> ::std::result::Result<Self, D::Error> {
        let raw = Value::deserialize(de)?;

        let mut table = match raw {
            Value::Table(t) => t,
            _ => {
                use serde::de::Error;
                return Err(D::Error::custom(
                    "A config file should always be a toml table",
                ));
            }
        };

        if is_legacy_format(&table) {
            warn!("It looks like you are using the legacy book.toml format.");
            warn!("We'll parse it for now, but you should probably convert to the new format.");
            warn!("See the mdbook documentation for more details");
            warn!("http://azerupi.github.io/mdBook/format/config.html");
            return Ok(Config::from_legacy(table));
        }

        let book: BookConfig = table.remove("book")
                                    .and_then(|value| value.try_into().ok())
                                    .unwrap_or_default();
        Ok(Config {
            book: book,
            rest: table,
        })
    }
}

fn is_legacy_format(table: &Table) -> bool {
    let top_level_items = ["title", "author", "authors"];

    top_level_items.iter().any(|key| table.contains_key(&key.to_string()))
}

/// Configuration options which are specific to the book and required for
/// loading it from disk.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(default, rename_all = "kebab-case")]
pub struct BookConfig {
    /// The book's title.
    pub title: Option<String>,
    /// The book's authors.
    pub authors: Vec<String>,
    /// An optional description for the book.
    pub description: Option<String>,
    /// Location of the book source relative to the book's root directory.
    pub src: PathBuf,
    /// Where to put built artefacts relative to the book's root directory.
    pub build_dir: PathBuf,
    /// Does this book support more than one language?
    pub multilingual: bool,
}

impl Default for BookConfig {
    fn default() -> BookConfig {
        BookConfig {
            title: None,
            authors: Vec::new(),
            description: None,
            src: PathBuf::from("src"),
            build_dir: PathBuf::from("book"),
            multilingual: false,
        }
    }
}

#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(default, rename_all = "kebab-case")]
pub struct HtmlConfig {
    pub theme: Option<PathBuf>,
    pub curly_quotes: bool,
    pub mathjax_support: bool,
    pub google_analytics: Option<String>,
    pub additional_css: Vec<PathBuf>,
    pub additional_js: Vec<PathBuf>,
    pub playpen: Playpen,
}

/// Configuration for tweaking how the the HTML renderer handles the playpen.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct Playpen {
    pub editor: PathBuf,
    pub editable: bool,
}

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

    const COMPLEX_CONFIG: &'static str = r#"
        [book]
        title = "Some Book"
        authors = ["Michael-F-Bryan <michaelfbryan@gmail.com>"]
        description = "A completely useless book"
        multilingual = true
        src = "source"
        build-dir = "outputs"

        [output.html]
        theme = "./themedir"
        curly-quotes = true
        google-analytics = "123456"
        additional-css = ["./foo/bar/baz.css"]

        [output.html.playpen]
        editable = true
        editor = "ace"
        "#;

    #[test]
    fn load_a_complex_config_file() {
        let src = COMPLEX_CONFIG;

        let book_should_be = BookConfig {
            title: Some(String::from("Some Book")),
            authors: vec![String::from("Michael-F-Bryan <michaelfbryan@gmail.com>")],
            description: Some(String::from("A completely useless book")),
            multilingual: true,
            src: PathBuf::from("source"),
            build_dir: PathBuf::from("outputs"),
            ..Default::default()
        };
        let playpen_should_be = Playpen {
            editable: true,
            editor: PathBuf::from("ace"),
        };
        let html_should_be = HtmlConfig {
            curly_quotes: true,
            google_analytics: Some(String::from("123456")),
            additional_css: vec![PathBuf::from("./foo/bar/baz.css")],
            theme: Some(PathBuf::from("./themedir")),
            playpen: playpen_should_be,
            ..Default::default()
        };

        let got = Config::from_str(src).unwrap();

        assert_eq!(got.book, book_should_be);
        assert_eq!(got.html_config().unwrap(), html_should_be);
    }

    #[test]
    fn load_arbitrary_output_type() {
        #[derive(Debug, Deserialize, PartialEq)]
        struct RandomOutput {
            foo: u32,
            bar: String,
            baz: Vec<bool>,
        }

        let src = r#"
        [output.random]
        foo = 5
        bar = "Hello World"
        baz = [true, true, false]
        "#;

        let should_be = RandomOutput {
            foo: 5,
            bar: String::from("Hello World"),
            baz: vec![true, true, false],
        };

        let cfg = Config::from_str(src).unwrap();
        let got: RandomOutput = cfg.get_deserialized("output.random").unwrap();

        assert_eq!(got, should_be);

        let baz: Vec<bool> = cfg.get_deserialized("output.random.baz").unwrap();
        let baz_should_be = vec![true, true, false];

        assert_eq!(baz, baz_should_be);
    }

    #[test]
    fn mutate_some_stuff() {
        // really this is just a sanity check to make sure the borrow checker
        // is happy...
        let src = COMPLEX_CONFIG;
        let mut config = Config::from_str(src).unwrap();
        let key = "output.html.playpen.editable";

        assert_eq!(config.get(key).unwrap(), &Value::Boolean(true));
        *config.get_mut(key).unwrap() = Value::Boolean(false);
        assert_eq!(config.get(key).unwrap(), &Value::Boolean(false));
    }

    /// The config file format has slightly changed (metadata stuff is now under
    /// the `book` table instead of being at the top level) so we're adding a
    /// **temporary** compatibility check. You should be able to still load the
    /// old format, emitting a warning.
    #[test]
    fn can_still_load_the_previous_format() {
        let src = r#"
        title = "mdBook Documentation"
        description = "Create book from markdown files. Like Gitbook but implemented in Rust"
        authors = ["Mathieu David"]
        source = "./source"

        [output.html]
        destination = "my-book" # the output files will be generated in `root/my-book` instead of `root/book`
        theme = "my-theme"
        curly-quotes = true
        google-analytics = "123456"
        additional-css = ["custom.css", "custom2.css"]
        additional-js = ["custom.js"]
        "#;

        let book_should_be = BookConfig {
            title: Some(String::from("mdBook Documentation")),
            description: Some(String::from(
                "Create book from markdown files. Like Gitbook but implemented in Rust",
            )),
            authors: vec![String::from("Mathieu David")],
            build_dir: PathBuf::from("my-book"),
            src: PathBuf::from("./source"),
            ..Default::default()
        };

        let html_should_be = HtmlConfig {
            theme: Some(PathBuf::from("my-theme")),
            curly_quotes: true,
            google_analytics: Some(String::from("123456")),
            additional_css: vec![PathBuf::from("custom.css"), PathBuf::from("custom2.css")],
            additional_js: vec![PathBuf::from("custom.js")],
            ..Default::default()
        };

        let got = Config::from_str(src).unwrap();
        assert_eq!(got.book, book_should_be);
        assert_eq!(got.html_config().unwrap(), html_should_be);
    }
}
<|endoftext|>


<commit_before>// spell-checker:ignore memmem
//! Implementations of digest functions, like md5 and sha1.
//!
//! The [`Digest`] trait represents the interface for providing inputs
//! to these digest functions and accessing the resulting hash. The
//! [`DigestWriter`] struct provides a wrapper around [`Digest`] that
//! implements the [`Write`] trait, for use in situations where calling
//! [`write`] would be useful.
extern crate digest;
extern crate md5;
extern crate sha1;
extern crate sha2;
extern crate sha3;

use std::io::Write;

use hex::encode;
#[cfg(windows)]
use memchr::memmem;

pub trait Digest {
    fn new() -> Self
    where
        Self: Sized;
    fn input(&mut self, input: &[u8]);
    fn result(&mut self, out: &mut [u8]);
    fn reset(&mut self);
    fn output_bits(&self) -> usize;
    fn output_bytes(&self) -> usize {
        (self.output_bits() + 7) / 8
    }
    fn result_str(&mut self) -> String {
        let mut buf: Vec<u8> = vec![0; self.output_bytes()];
        self.result(&mut buf);
        encode(buf)
    }
}

impl Digest for blake2b_simd::State {
    fn new() -> Self {
        Self::new()
    }

    fn input(&mut self, input: &[u8]) {
        self.update(input);
    }

    fn result(&mut self, out: &mut [u8]) {
        let hash_result = &self.finalize();
        out.copy_from_slice(hash_result.as_bytes());
    }

    fn reset(&mut self) {
        *self = Self::new();
    }

    fn output_bits(&self) -> usize {
        512
    }
}

impl Digest for blake3::Hasher {
    fn new() -> Self {
        Self::new()
    }

    fn input(&mut self, input: &[u8]) {
        self.update(input);
    }

    fn result(&mut self, out: &mut [u8]) {
        let hash_result = &self.finalize();
        out.copy_from_slice(hash_result.as_bytes());
    }

    fn reset(&mut self) {
        *self = Self::new();
    }

    fn output_bits(&self) -> usize {
        256
    }
}

// Implements the Digest trait for sha2 / sha3 algorithms with fixed output
macro_rules! impl_digest_rustcryptocommon {
    ($type: ty, $size: expr) => {
        impl Digest for $type {
            fn new() -> Self {
                Self::default()
            }

            fn input(&mut self, input: &[u8]) {
                digest::Digest::update(self, input);
            }

            fn result(&mut self, out: &mut [u8]) {
                digest::Digest::finalize_into_reset(self, out.into());
            }

            fn reset(&mut self) {
                *self = Self::new();
            }

            fn output_bits(&self) -> usize {
                $size
            }
        }
    };
}

// Implements the Digest trait for sha2 / sha3 algorithms with variable output
macro_rules! impl_digest_shake {
    ($type: ty) => {
        impl Digest for $type {
            fn new() -> Self {
                Self::default()
            }

            fn input(&mut self, input: &[u8]) {
                digest::Update::update(self, input);
            }

            fn result(&mut self, out: &mut [u8]) {
                digest::ExtendableOutputReset::finalize_xof_reset_into(self, out);
            }

            fn reset(&mut self) {
                *self = Self::new();
            }

            fn output_bits(&self) -> usize {
                0
            }
        }
    };
}

impl_digest_rustcryptocommon!(md5::Md5, 128);
impl_digest_rustcryptocommon!(sha1::Sha1, 160);
impl_digest_rustcryptocommon!(sha2::Sha224, 224);
impl_digest_rustcryptocommon!(sha2::Sha256, 256);
impl_digest_rustcryptocommon!(sha2::Sha384, 384);
impl_digest_rustcryptocommon!(sha2::Sha512, 512);

impl_digest_rustcryptocommon!(sha3::Sha3_224, 224);
impl_digest_rustcryptocommon!(sha3::Sha3_256, 256);
impl_digest_rustcryptocommon!(sha3::Sha3_384, 384);
impl_digest_rustcryptocommon!(sha3::Sha3_512, 512);
impl_digest_shake!(sha3::Shake128);
impl_digest_shake!(sha3::Shake256);

/// A struct that writes to a digest.
///
/// This struct wraps a [`Digest`] and provides a [`Write`]
/// implementation that passes input bytes directly to the
/// [`Digest::input`].
///
/// On Windows, if `binary` is `false`, then the [`write`]
/// implementation replaces instances of "\r\n" with "\n" before passing
/// the input bytes to the [`digest`].
pub struct DigestWriter<'a> {
    digest: &'a mut Box<dyn Digest>,

    /// Whether to write to the digest in binary mode or text mode on Windows.
    ///
    /// If this is `false`, then instances of "\r\n" are replaced with
    /// "\n" before passing input bytes to the [`digest`].
    #[allow(dead_code)]
    binary: bool,

    /// Whether the previous
    #[allow(dead_code)]
    was_last_character_carriage_return: bool,
    // TODO These are dead code only on non-Windows operating systems.
    // It might be better to use a `#[cfg(windows)]` guard here.
}

impl<'a> DigestWriter<'a> {
    pub fn new(digest: &'a mut Box<dyn Digest>, binary: bool) -> DigestWriter {
        let was_last_character_carriage_return = false;
        DigestWriter {
            digest,
            binary,
            was_last_character_carriage_return,
        }
    }

    pub fn finalize(&mut self) -> bool {
        if self.was_last_character_carriage_return {
            self.digest.input(&[b'\r']);
            true
        } else {
            false
        }
    }
}

impl<'a> Write for DigestWriter<'a> {
    #[cfg(not(windows))]
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.digest.input(buf);
        Ok(buf.len())
    }

    #[cfg(windows)]
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        if self.binary {
            self.digest.input(buf);
            return Ok(buf.len());
        }

        // The remaining code handles Windows text mode, where we must
        // replace each occurrence of "\r\n" with "\n".
        //
        // First, if the last character written was "\r" and the first
        // character in the current buffer to write is not "\n", then we
        // need to write the "\r" that we buffered from the previous
        // call to `write()`.
        let n = buf.len();
        if self.was_last_character_carriage_return && n > 0 && buf[0] != b'\n' {
            self.digest.input(&[b'\r']);
        }

        // Next, find all occurrences of "\r\n", inputting the slice
        // just before the "\n" in the previous instance of "\r\n" and
        // the beginning of this "\r\n".
        let mut i_prev = 0;
        for i in memmem::find_iter(buf, b"\r\n") {
            self.digest.input(&buf[i_prev..i]);
            i_prev = i + 1;
        }

        // Finally, check whether the last character is "\r". If so,
        // buffer it until we know that the next character is not "\n",
        // which can only be known on the next call to `write()`.
        //
        // This all assumes that `write()` will be called on adjacent
        // blocks of the input.
        if n > 0 && buf[n - 1] == b'\r' {
            self.was_last_character_carriage_return = true;
            self.digest.input(&buf[i_prev..n - 1]);
        } else {
            self.was_last_character_carriage_return = false;
            self.digest.input(&buf[i_prev..n]);
        }

        // Even though we dropped a "\r" for each "\r\n" we found, we
        // still report the number of bytes written as `n`. This is
        // because the meaning of the returned number is supposed to be
        // the number of bytes consumed by the writer, so that if the
        // calling code were calling `write()` in a loop, it would know
        // where the next contiguous slice of the buffer starts.
        Ok(n)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {

    /// Test for replacing a "\r\n" sequence with "\n" when the "\r" is
    /// at the end of one block and the "\n" is at the beginning of the
    /// next block, when reading in blocks.
    #[cfg(windows)]
    #[test]
    fn test_crlf_across_blocks() {
        use std::io::Write;

        use crate::digest::Digest;
        use crate::digest::DigestWriter;

        // Writing "\r" in one call to `write()`, and then "\n" in another.
        let mut digest = Box::new(md5::Context::new()) as Box<dyn Digest>;
        let mut writer_crlf = DigestWriter::new(&mut digest, false);
        writer_crlf.write_all(&[b'\r']).unwrap();
        writer_crlf.write_all(&[b'\n']).unwrap();
        writer_crlf.finalize();
        let result_crlf = digest.result_str();

        // We expect "\r\n" to be replaced with "\n" in text mode on Windows.
        let mut digest = Box::new(md5::Context::new()) as Box<dyn Digest>;
        let mut writer_lf = DigestWriter::new(&mut digest, false);
        writer_lf.write_all(&[b'\n']).unwrap();
        writer_lf.finalize();
        let result_lf = digest.result_str();

        assert_eq!(result_crlf, result_lf);
    }
}
<commit_msg>hashsum: pass spell checker<commit_after>// spell-checker:ignore memmem
//! Implementations of digest functions, like md5 and sha1.
//!
//! The [`Digest`] trait represents the interface for providing inputs
//! to these digest functions and accessing the resulting hash. The
//! [`DigestWriter`] struct provides a wrapper around [`Digest`] that
//! implements the [`Write`] trait, for use in situations where calling
//! [`write`] would be useful.
extern crate digest;
extern crate md5;
extern crate sha1;
extern crate sha2;
extern crate sha3;

use std::io::Write;

use hex::encode;
#[cfg(windows)]
use memchr::memmem;

pub trait Digest {
    fn new() -> Self
    where
        Self: Sized;
    fn input(&mut self, input: &[u8]);
    fn result(&mut self, out: &mut [u8]);
    fn reset(&mut self);
    fn output_bits(&self) -> usize;
    fn output_bytes(&self) -> usize {
        (self.output_bits() + 7) / 8
    }
    fn result_str(&mut self) -> String {
        let mut buf: Vec<u8> = vec![0; self.output_bytes()];
        self.result(&mut buf);
        encode(buf)
    }
}

impl Digest for blake2b_simd::State {
    fn new() -> Self {
        Self::new()
    }

    fn input(&mut self, input: &[u8]) {
        self.update(input);
    }

    fn result(&mut self, out: &mut [u8]) {
        let hash_result = &self.finalize();
        out.copy_from_slice(hash_result.as_bytes());
    }

    fn reset(&mut self) {
        *self = Self::new();
    }

    fn output_bits(&self) -> usize {
        512
    }
}

impl Digest for blake3::Hasher {
    fn new() -> Self {
        Self::new()
    }

    fn input(&mut self, input: &[u8]) {
        self.update(input);
    }

    fn result(&mut self, out: &mut [u8]) {
        let hash_result = &self.finalize();
        out.copy_from_slice(hash_result.as_bytes());
    }

    fn reset(&mut self) {
        *self = Self::new();
    }

    fn output_bits(&self) -> usize {
        256
    }
}

// Implements the Digest trait for sha2 / sha3 algorithms with fixed output
macro_rules! impl_digest_common {
    ($type: ty, $size: expr) => {
        impl Digest for $type {
            fn new() -> Self {
                Self::default()
            }

            fn input(&mut self, input: &[u8]) {
                digest::Digest::update(self, input);
            }

            fn result(&mut self, out: &mut [u8]) {
                digest::Digest::finalize_into_reset(self, out.into());
            }

            fn reset(&mut self) {
                *self = Self::new();
            }

            fn output_bits(&self) -> usize {
                $size
            }
        }
    };
}

// Implements the Digest trait for sha2 / sha3 algorithms with variable output
macro_rules! impl_digest_shake {
    ($type: ty) => {
        impl Digest for $type {
            fn new() -> Self {
                Self::default()
            }

            fn input(&mut self, input: &[u8]) {
                digest::Update::update(self, input);
            }

            fn result(&mut self, out: &mut [u8]) {
                digest::ExtendableOutputReset::finalize_xof_reset_into(self, out);
            }

            fn reset(&mut self) {
                *self = Self::new();
            }

            fn output_bits(&self) -> usize {
                0
            }
        }
    };
}

impl_digest_common!(md5::Md5, 128);
impl_digest_common!(sha1::Sha1, 160);
impl_digest_common!(sha2::Sha224, 224);
impl_digest_common!(sha2::Sha256, 256);
impl_digest_common!(sha2::Sha384, 384);
impl_digest_common!(sha2::Sha512, 512);

impl_digest_common!(sha3::Sha3_224, 224);
impl_digest_common!(sha3::Sha3_256, 256);
impl_digest_common!(sha3::Sha3_384, 384);
impl_digest_common!(sha3::Sha3_512, 512);
impl_digest_shake!(sha3::Shake128);
impl_digest_shake!(sha3::Shake256);

/// A struct that writes to a digest.
///
/// This struct wraps a [`Digest`] and provides a [`Write`]
/// implementation that passes input bytes directly to the
/// [`Digest::input`].
///
/// On Windows, if `binary` is `false`, then the [`write`]
/// implementation replaces instances of "\r\n" with "\n" before passing
/// the input bytes to the [`digest`].
pub struct DigestWriter<'a> {
    digest: &'a mut Box<dyn Digest>,

    /// Whether to write to the digest in binary mode or text mode on Windows.
    ///
    /// If this is `false`, then instances of "\r\n" are replaced with
    /// "\n" before passing input bytes to the [`digest`].
    #[allow(dead_code)]
    binary: bool,

    /// Whether the previous
    #[allow(dead_code)]
    was_last_character_carriage_return: bool,
    // TODO These are dead code only on non-Windows operating systems.
    // It might be better to use a `#[cfg(windows)]` guard here.
}

impl<'a> DigestWriter<'a> {
    pub fn new(digest: &'a mut Box<dyn Digest>, binary: bool) -> DigestWriter {
        let was_last_character_carriage_return = false;
        DigestWriter {
            digest,
            binary,
            was_last_character_carriage_return,
        }
    }

    pub fn finalize(&mut self) -> bool {
        if self.was_last_character_carriage_return {
            self.digest.input(&[b'\r']);
            true
        } else {
            false
        }
    }
}

impl<'a> Write for DigestWriter<'a> {
    #[cfg(not(windows))]
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.digest.input(buf);
        Ok(buf.len())
    }

    #[cfg(windows)]
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        if self.binary {
            self.digest.input(buf);
            return Ok(buf.len());
        }

        // The remaining code handles Windows text mode, where we must
        // replace each occurrence of "\r\n" with "\n".
        //
        // First, if the last character written was "\r" and the first
        // character in the current buffer to write is not "\n", then we
        // need to write the "\r" that we buffered from the previous
        // call to `write()`.
        let n = buf.len();
        if self.was_last_character_carriage_return && n > 0 && buf[0] != b'\n' {
            self.digest.input(&[b'\r']);
        }

        // Next, find all occurrences of "\r\n", inputting the slice
        // just before the "\n" in the previous instance of "\r\n" and
        // the beginning of this "\r\n".
        let mut i_prev = 0;
        for i in memmem::find_iter(buf, b"\r\n") {
            self.digest.input(&buf[i_prev..i]);
            i_prev = i + 1;
        }

        // Finally, check whether the last character is "\r". If so,
        // buffer it until we know that the next character is not "\n",
        // which can only be known on the next call to `write()`.
        //
        // This all assumes that `write()` will be called on adjacent
        // blocks of the input.
        if n > 0 && buf[n - 1] == b'\r' {
            self.was_last_character_carriage_return = true;
            self.digest.input(&buf[i_prev..n - 1]);
        } else {
            self.was_last_character_carriage_return = false;
            self.digest.input(&buf[i_prev..n]);
        }

        // Even though we dropped a "\r" for each "\r\n" we found, we
        // still report the number of bytes written as `n`. This is
        // because the meaning of the returned number is supposed to be
        // the number of bytes consumed by the writer, so that if the
        // calling code were calling `write()` in a loop, it would know
        // where the next contiguous slice of the buffer starts.
        Ok(n)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {

    /// Test for replacing a "\r\n" sequence with "\n" when the "\r" is
    /// at the end of one block and the "\n" is at the beginning of the
    /// next block, when reading in blocks.
    #[cfg(windows)]
    #[test]
    fn test_crlf_across_blocks() {
        use std::io::Write;

        use crate::digest::Digest;
        use crate::digest::DigestWriter;

        // Writing "\r" in one call to `write()`, and then "\n" in another.
        let mut digest = Box::new(md5::Context::new()) as Box<dyn Digest>;
        let mut writer_crlf = DigestWriter::new(&mut digest, false);
        writer_crlf.write_all(&[b'\r']).unwrap();
        writer_crlf.write_all(&[b'\n']).unwrap();
        writer_crlf.finalize();
        let result_crlf = digest.result_str();

        // We expect "\r\n" to be replaced with "\n" in text mode on Windows.
        let mut digest = Box::new(md5::Context::new()) as Box<dyn Digest>;
        let mut writer_lf = DigestWriter::new(&mut digest, false);
        writer_lf.write_all(&[b'\n']).unwrap();
        writer_lf.finalize();
        let result_lf = digest.result_str();

        assert_eq!(result_crlf, result_lf);
    }
}
<|endoftext|>


<commit_before>use wire::handler::{PeerId, PeerAction, ServerHandler};
use metainfo::MetaInfo;
use metainfo::SHA1Hash20b;
use file::PartialFile;
use std::collections::HashMap;
use wire::data::PeerMsg;
use std::time::SystemTime;
use file::{PartialFileTrait, PeerFile};

pub struct PeerState {
    has_handshake: bool,
    disconnected: bool,
    peer_choking: bool,
    peer_interested: bool,
    am_choking: bool,
    am_interested: bool,
    last_keepalive: SystemTime,
    file: PeerFile,
}

impl PeerState {
    pub fn new(len: usize) -> PeerState {
        PeerState {
            has_handshake: false,
            disconnected: false,
            peer_choking: true,
            peer_interested: false,
            am_choking: true,
            am_interested: false,
            last_keepalive: SystemTime::now(),
            file: PeerFile::new(len),
        }
    }
}

pub struct PeerServer {
    peers: HashMap<PeerId, PeerState>,
    hash: SHA1Hash20b,
    our_peer_id: String,
    partial_file: PartialFile,
}

const PROTOCOL_ID: &'static str = "rustorrent-beta";

impl ServerHandler for PeerServer {
    fn new(metainfo: MetaInfo, hash: SHA1Hash20b, our_peer_id: &str) -> Self {
        let partial_file = PartialFile::new(&metainfo.info);

        PeerServer {
            peers: HashMap::new(),
            hash: hash,
            our_peer_id: our_peer_id.to_string(),
            partial_file: partial_file,
        }
    }

    fn on_peer_connect(&mut self, id: PeerId) -> PeerAction {
        let handshake = PeerMsg::handshake(PROTOCOL_ID.to_string(),
                                           self.our_peer_id.to_string(),
                                           &self.hash);
        PeerAction::SendMessages(vec![handshake])
    }

    fn on_message_receive(&mut self, id: PeerId, msg: PeerMsg) -> PeerAction {
        {
            let peer = match self.peers.get_mut(&id) {
                Some(peer) => peer,
                None => return PeerAction::Nothing,
            };

            if peer.disconnected {
                return PeerAction::Nothing;
            }

            if !peer.has_handshake {
                match msg {
                    PeerMsg::HandShake(_, ref their_hash, _) => {
                        if their_hash == &self.hash {
                            peer.has_handshake = true;
                        } else {
                            peer.disconnected = false;
                        }
                        return PeerAction::Nothing;
                    }
                    _ => peer.disconnected = true,
                }
            }
        }

        let mut outgoing_msgs = Vec::new();
        let mut handled = true;

        {
            let peer = match self.peers.get_mut(&id) {
                Some(peer) => peer,
                None => return PeerAction::Nothing,
            };

            // messages that mutate the peer
            match msg {
                PeerMsg::HandShake(..) => {}
                PeerMsg::KeepAlive => peer.last_keepalive = SystemTime::now(),
                PeerMsg::Choke => {
                    peer.peer_choking = true;
                }
                PeerMsg::Unchoke => {
                    peer.peer_choking = false;
                }
                PeerMsg::Interested => {
                    peer.peer_interested = true;
                }
                PeerMsg::NotInterested => {
                    peer.peer_interested = false;
                }
                PeerMsg::Have(index) => peer.file.set(index as usize, true),
                PeerMsg::Bitfield(_) => {}
                _ => handled = true,
            };
        }

        // messages that don't need to mutate peer
        let choking = {
            match self.peers.get(&id) {
                Some(peer) => peer.am_choking,
                None => return PeerAction::Nothing,
            }
        };

        //
        if !choking && !handled {
            match msg {
                PeerMsg::Request(index, begin, offset) => {
                    let response = self._get_piece_from_req(index as usize, begin, offset);
                    if response.is_some() {
                        outgoing_msgs.push(response.unwrap());
                    }
                }
                PeerMsg::Piece(index, begin, block) => {
                    self.partial_file.add_piece(index as usize, begin as usize, block);
                }
                PeerMsg::Cancel(..) => {}
                PeerMsg::Port(_) => {}
                _ => handled = true,
            }
        }

        if outgoing_msgs.len() > 0 {
            PeerAction::SendMessages(outgoing_msgs)
        } else {
            PeerAction::Nothing
        }
    }

    fn on_peer_disconnect(&mut self, id: PeerId) -> PeerAction {
        unimplemented!();
    }

    fn on_loop(&mut self) -> PeerAction {
        unimplemented!();
    }
}

impl PeerServer {
    fn _get_piece_from_req(&mut self, index: usize, begin: u32, offset: u32) -> Option<PeerMsg> {
        if self.partial_file.has_piece(index as usize) {
            let piece = self.partial_file.get_piece_mut(index as usize);
            return match piece.get_offset(begin as usize, offset as usize) {
                Some(piece_data) => {
                    let msg = PeerMsg::Piece(begin, offset, Vec::from(piece_data));
                    Some(msg)
                }
                _ => None,
            };
        }
        None
    }
}
<commit_msg>Add functionality for removing timedout peers<commit_after>use wire::handler::{PeerId, PeerAction, ServerHandler};
use metainfo::MetaInfo;
use metainfo::SHA1Hash20b;
use file::PartialFile;
use std::collections::HashMap;
use wire::data::PeerMsg;
use std::time::SystemTime;
use file::{PartialFileTrait, PeerFile};

const TIMEOUT_SECONDS: u64 = 60 * 5;

pub struct PeerState {
    has_handshake: bool,
    disconnected: bool,
    peer_choking: bool,
    peer_interested: bool,
    am_choking: bool,
    am_interested: bool,
    last_keepalive: SystemTime,
    file: PeerFile,
}

impl PeerState {
    pub fn new(len: usize) -> PeerState {
        PeerState {
            has_handshake: false,
            disconnected: false,
            peer_choking: true,
            peer_interested: false,
            am_choking: true,
            am_interested: false,
            last_keepalive: SystemTime::now(),
            file: PeerFile::new(len),
        }
    }
}

pub struct PeerServer {
    peers: HashMap<PeerId, PeerState>,
    hash: SHA1Hash20b,
    our_peer_id: String,
    partial_file: PartialFile,
}

const PROTOCOL_ID: &'static str = "rustorrent-beta";

impl ServerHandler for PeerServer {
    fn new(metainfo: MetaInfo, hash: SHA1Hash20b, our_peer_id: &str) -> Self {
        let partial_file = PartialFile::new(&metainfo.info);

        PeerServer {
            peers: HashMap::new(),
            hash: hash,
            our_peer_id: our_peer_id.to_string(),
            partial_file: partial_file,
        }
    }

    fn on_peer_connect(&mut self, id: PeerId) -> PeerAction {
        let handshake = PeerMsg::handshake(PROTOCOL_ID.to_string(),
                                           self.our_peer_id.to_string(),
                                           &self.hash);
        PeerAction::SendMessages(vec![handshake])
    }

    fn on_message_receive(&mut self, id: PeerId, msg: PeerMsg) -> PeerAction {
        {
            let peer = match self.peers.get_mut(&id) {
                Some(peer) => peer,
                None => return PeerAction::Nothing,
            };

            if peer.disconnected {
                return PeerAction::Nothing;
            }

            if !peer.has_handshake {
                match msg {
                    PeerMsg::HandShake(_, ref their_hash, _) => {
                        if their_hash == &self.hash {
                            peer.has_handshake = true;
                        } else {
                            peer.disconnected = false;
                        }
                        return PeerAction::Nothing;
                    }
                    _ => peer.disconnected = true,
                }
            }
        }

        let mut outgoing_msgs = Vec::new();
        let mut handled = true;

        {
            let peer = match self.peers.get_mut(&id) {
                Some(peer) => peer,
                None => return PeerAction::Nothing,
            };

            // messages that mutate the peer
            match msg {
                PeerMsg::HandShake(..) => {}
                PeerMsg::KeepAlive => peer.last_keepalive = SystemTime::now(),
                PeerMsg::Choke => {
                    peer.peer_choking = true;
                }
                PeerMsg::Unchoke => {
                    peer.peer_choking = false;
                }
                PeerMsg::Interested => {
                    peer.peer_interested = true;
                }
                PeerMsg::NotInterested => {
                    peer.peer_interested = false;
                }
                PeerMsg::Have(index) => peer.file.set(index as usize, true),
                PeerMsg::Bitfield(_) => {}
                _ => handled = false,
            };
        }

        // messages that don't need to mutate peer
        let choking = {
            match self.peers.get(&id) {
                Some(peer) => peer.am_choking,
                None => return PeerAction::Nothing,
            }
        };

        if !choking && !handled {
            match msg {
                PeerMsg::Request(index, begin, offset) => {
                    let response = self._get_piece_from_req(index as usize, begin, offset);
                    if response.is_some() {
                        outgoing_msgs.push(response.unwrap());
                    }
                }
                PeerMsg::Piece(index, begin, block) => {
                    self.partial_file.add_piece(index as usize, begin as usize, block);
                }
                PeerMsg::Cancel(..) => {}
                PeerMsg::Port(_) => {}
                _ => handled = true,
            }
        }

        if outgoing_msgs.len() > 0 {
            PeerAction::SendMessages(outgoing_msgs)
        } else {
            PeerAction::Nothing
        }
    }

    fn on_peer_disconnect(&mut self, id: PeerId) -> PeerAction {
        self.peers.remove(&id);
        PeerAction::Nothing
    }

    //remove peers that have no replied in five minutes
    fn on_loop(&mut self) -> PeerAction {
        self._remove_old_peers();
        PeerAction::Nothing
    }
}

impl PeerServer {
    fn _remove_old_peers(&mut self) {

        let for_removal = self._get_timeout_ids();

        for id in for_removal {
            self.peers.remove(&id);
        }

    }

    fn _get_timeout_ids(&self) -> Vec<PeerId> {
        let mut for_removal = Vec::new();
            for (&id, peer) in &self.peers {
                match peer.last_keepalive.elapsed() {
                    Ok(duration) => {
                        if duration.as_secs() > TIMEOUT_SECONDS {
                            for_removal.push(id);
                        }
                    },
                    _ => ()
                }
            }
            for_removal
    }

    fn _get_piece_from_req(&mut self, index: usize, begin: u32, offset: u32) -> Option<PeerMsg> {
        if self.partial_file.has_piece(index as usize) {
            let piece = self.partial_file.get_piece_mut(index as usize);
            return match piece.get_offset(begin as usize, offset as usize) {
                Some(piece_data) => {
                    let msg = PeerMsg::Piece(begin, offset, Vec::from(piece_data));
                    Some(msg)
                }
                _ => None,
            };
        }
        None
    }
}
<|endoftext|>