<commit_before>// Copyright 2017 The xi-editor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Management of styles.

use std::collections::{HashMap, HashSet};
use std::ffi::OsStr;
use std::fs;
use std::iter::FromIterator;
use std::path::{Path, PathBuf};

use serde_json::{self, Value};
use syntect::dumps::{dump_to_file, from_dump_file};
use syntect::highlighting::StyleModifier as SynStyleModifier;
use syntect::highlighting::{Color, Highlighter, Theme, ThemeSet};
use syntect::LoadingError;

pub use syntect::highlighting::ThemeSettings;

pub const N_RESERVED_STYLES: usize = 8;
const SYNTAX_PRIORITY_DEFAULT: u16 = 200;
const SYNTAX_PRIORITY_LOWEST: u16 = 0;
pub const DEFAULT_THEME: &str = "InspiredGitHub";

#[derive(Clone, PartialEq, Eq, Default, Hash, Debug, Serialize, Deserialize)]
/// A mergeable style. All values except priority are optional.
///
/// Note: A `None` value represents the absense of preference; in the case of
/// boolean options, `Some(false)` means that this style will override a lower
/// priority value in the same field.
pub struct Style {
    /// The priority of this style, in the range (0, 1000). Used to resolve
    /// conflicting fields when merging styles. The higher priority wins.
    #[serde(skip_serializing)]
    pub priority: u16,
    /// The foreground text color, in ARGB.
    pub fg_color: Option<u32>,
    /// The background text color, in ARGB.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bg_color: Option<u32>,
    /// The font-weight, in the range 100-900, interpreted like the CSS
    /// font-weight property.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub weight: Option<u16>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub underline: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub italic: Option<bool>,
}

impl Style {
    /// Creates a new `Style` by converting from a `Syntect::StyleModifier`.
    pub fn from_syntect_style_mod(style: &SynStyleModifier) -> Self {
        let font_style = style.font_style.map(|s| s.bits()).unwrap_or_default();
        let weight = if (font_style & 1) != 0 { Some(700) } else { None };
        let underline = if (font_style & 2) != 0 { Some(true) } else { None };
        let italic = if (font_style & 4) != 0 { Some(true) } else { None };

        Self::new(
            SYNTAX_PRIORITY_DEFAULT,
            style.foreground.map(Self::rgba_from_syntect_color),
            style.background.map(Self::rgba_from_syntect_color),
            weight,
            underline,
            italic,
        )
    }

    pub fn new<O32, O16, OB>(
        priority: u16,
        fg_color: O32,
        bg_color: O32,
        weight: O16,
        underline: OB,
        italic: OB,
    ) -> Self
    where
        O32: Into<Option<u32>>,
        O16: Into<Option<u16>>,
        OB: Into<Option<bool>>,
    {
        assert!(priority <= 1000);
        Style {
            priority,
            fg_color: fg_color.into(),
            bg_color: bg_color.into(),
            weight: weight.into(),
            underline: underline.into(),
            italic: italic.into(),
        }
    }

    /// Returns the default style for the given `Theme`.
    pub fn default_for_theme(theme: &Theme) -> Self {
        let fg = theme.settings.foreground.unwrap_or(Color::BLACK);
        Style::new(
            SYNTAX_PRIORITY_LOWEST,
            Some(Self::rgba_from_syntect_color(fg)),
            None,
            None,
            None,
            None,
        )
    }

    /// Creates a new style by combining attributes of `self` and `other`.
    /// If both styles define an attribute, the highest priority wins; `other`
    /// wins in the case of a tie.
    ///
    /// Note: when merging multiple styles, apply them in increasing priority.
    pub fn merge(&self, other: &Style) -> Style {
        let (p1, p2) = if self.priority > other.priority { (self, other) } else { (other, self) };

        Style::new(
            p1.priority,
            p1.fg_color.or(p2.fg_color),
            p1.bg_color.or(p2.bg_color),
            p1.weight.or(p2.weight),
            p1.underline.or(p2.underline),
            p1.italic.or(p2.italic),
        )
    }

    /// Encode this `Style`, setting the `id` property.
    ///
    /// Note: this should only be used when sending the `def_style` RPC.
    pub fn to_json(&self, id: usize) -> Value {
        let mut as_val = serde_json::to_value(self).expect("failed to encode style");
        as_val["id"] = id.into();
        as_val
    }

    fn rgba_from_syntect_color(color: Color) -> u32 {
        let Color { r, g, b, a } = color;
        ((a as u32) << 24) | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32)
    }
}

/// A map from styles to client identifiers for a given `Theme`.
pub struct ThemeStyleMap {
    themes: ThemeSet,
    theme_name: String,
    theme: Theme,
    default_style: Style,
    map: HashMap<Style, usize>,

    // It's not obvious we actually have to store the style, we seem to only need it
    // as the key in the map.
    styles: Vec<Style>,
    themes_dir: Option<PathBuf>,
    cache_dir: Option<PathBuf>,
    caching_enabled: bool,

    // Maintaining all theme paths for comparison on an fs event.
    state: HashSet<PathBuf>,
}

impl ThemeStyleMap {
    pub fn new(themes_dir: Option<PathBuf>) -> ThemeStyleMap {
        let themes = ThemeSet::load_defaults();
        let theme_name = DEFAULT_THEME.to_owned();
        let theme = themes.themes.get(&theme_name).expect("missing theme").to_owned();
        let default_style = Style::default_for_theme(&theme);
        let cache_dir = None;
        let caching_enabled = true;

        ThemeStyleMap {
            themes,
            theme_name,
            theme,
            default_style,
            map: HashMap::new(),
            styles: Vec::new(),
            themes_dir,
            cache_dir,
            caching_enabled,
            state: HashSet::new(),
        }
    }

    pub fn get_default_style(&self) -> &Style {
        &self.default_style
    }

    pub fn get_highlighter(&self) -> Highlighter {
        Highlighter::new(&self.theme)
    }

    pub fn get_theme_name(&self) -> &str {
        &self.theme_name
    }

    pub fn get_theme_settings(&self) -> &ThemeSettings {
        &self.theme.settings
    }

    pub fn get_theme_names(&self) -> Vec<String> {
        self.themes.themes.keys().cloned().collect()
    }

    pub fn contains_theme(&self, k: &str) -> bool {
        self.themes.themes.contains_key(k)
    }

    pub fn set_theme(&mut self, theme_name: &str) -> Result<(), &'static str> {
        if let Some(new_theme) = self.themes.themes.get(theme_name) {
            self.theme = new_theme.to_owned();
            self.theme_name = theme_name.to_owned();
            self.default_style = Style::default_for_theme(&self.theme);
            self.map = HashMap::new();
            self.styles = Vec::new();
            Ok(())
        } else {
            Err("unknown theme")
        }
    }

    pub fn merge_with_default(&self, style: &Style) -> Style {
        self.default_style.merge(style)
    }

    pub fn lookup(&self, style: &Style) -> Option<usize> {
        self.map.get(style).cloned()
    }

    pub fn add(&mut self, style: &Style) -> usize {
        let result = self.styles.len() + N_RESERVED_STYLES;
        self.map.insert(style.clone(), result);
        self.styles.push(style.clone());
        result
    }

    /// Delete key and the corresponding dump file from the themes map.
    pub(crate) fn remove_theme(&mut self, path: &Path) -> Option<String> {
        validate_theme_file(path).ok()?;

        let theme_name = path.file_stem().and_then(OsStr::to_str)?;
        self.themes.themes.remove(theme_name);
        self.state.remove(path);

        let dump_p = self.get_dump_path(theme_name)?;
        if dump_p.exists() {
            let _ = fs::remove_file(dump_p);
        }

        Some(theme_name.to_string())
    }

    /// Load all themes inside the given directory.
    pub(crate) fn load_theme_dir(&mut self) {
        if let Some(themes_dir) = self.themes_dir.clone() {
            match ThemeSet::discover_theme_paths(themes_dir) {
                Ok(themes) => {
                    self.caching_enabled = self.caching_enabled && self.init_cache_dir();

                    for theme_p in &themes {
                        match self.try_load_from_dump(theme_p) {
                            Some((k, v)) => {
                                self.insert_to_map(k, v, theme_p);
                            }
                            None => {
                                let _ = self.load_theme(theme_p);
                            }
                        }
                    }
                }
                Err(e) => error!("Error loading themes dir: {:?}", e),
            }
        }
    }

    /// A wrapper around `from_dump_file`
    /// to validate the state of dump file.
    /// Invalidates if mod time of dump is less
    /// than the original one.
    fn try_load_from_dump(&self, theme_p: &Path) -> Option<(String, Theme)> {
        if !self.caching_enabled {
            return None;
        }

        let theme_name = theme_p.file_stem().and_then(OsStr::to_str)?;

        let dump_p = self.get_dump_path(theme_name)?;

        if !&dump_p.exists() {
            return None;
        }

        //NOTE: `try_load_from_dump` will return `None` if the file at
        //`dump_p` or `theme_p` is deleted before the execution of this fn.
        let mod_t = fs::metadata(&dump_p).and_then(|md| md.modified()).ok()?;
        let mod_t_orig = fs::metadata(theme_p).and_then(|md| md.modified()).ok()?;

        if mod_t >= mod_t_orig {
            from_dump_file(&dump_p).ok().map(|t| (theme_name.to_owned(), t))
        } else {
            // Delete dump file
            let _ = fs::remove_file(&dump_p);
            None
        }
    }

    /// Loads theme using syntect's `get_theme` fn to our `theme` map.
    /// Stores binary dump in a file with `tmdump` extension, only if
    /// caching is enabled.
    pub(crate) fn load_theme(&mut self, theme_p: &Path) -> Result<String, LoadingError> {
        validate_theme_file(theme_p)?;
        let theme = ThemeSet::get_theme(theme_p)?;
        let theme_name =
            theme_p.file_stem().and_then(OsStr::to_str).ok_or(LoadingError::BadPath)?;

        if self.caching_enabled {
            if let Some(dump_p) = self.get_dump_path(theme_name) {
                let _ = dump_to_file(&theme, dump_p);
            }
        }
        self.insert_to_map(theme_name.to_owned(), theme, theme_p);
        Ok(theme_name.to_owned())
    }

    fn insert_to_map(&mut self, k: String, v: Theme, p: &Path) {
        self.themes.themes.insert(k, v);

        //Maintain a record for future syncing
        self.state.insert(p.to_path_buf());
    }

    /// Returns dump's path corresponding to the given theme name.
    fn get_dump_path(&self, theme_name: &str) -> Option<PathBuf> {
        self.cache_dir.as_ref().map(|p| p.join(theme_name).with_extension("tmdump"))
    }

    /// Compare the stored file paths in `self.state`
    /// to the present ones.
    pub(crate) fn sync_dir(&mut self, dir: Option<&Path>) {
        if let Some(themes_dir) = dir {
            if let Ok(paths) = ThemeSet::discover_theme_paths(themes_dir) {
                let current_state = HashSet::from_iter(paths.into_iter());
                let maintained_state = self.state.clone();

                let to_insert = current_state.difference(&maintained_state);
                for path in to_insert {
                    let _ = self.load_theme(path);
                }

                let to_remove = maintained_state.difference(&current_state);
                for path in to_remove {
                    self.remove_theme(path);
                }
            }
        }
    }

    /// Creates the cache dir returns true
    /// if it is successfully initialized or
    /// already exists.
    fn init_cache_dir(&mut self) -> bool {
        self.cache_dir = self.themes_dir.clone().map(|p| p.join("cache"));

        if let Some(ref p) = self.cache_dir {
            if p.exists() {
                return true;
            }
            fs::DirBuilder::new().create(&p).is_ok()
        } else {
            false
        }
    }
}

/// Used to remove files with extension other than `tmTheme`.
fn validate_theme_file(path: &Path) -> Result<bool, LoadingError> {
    path.extension().map(|e| e != "tmTheme").ok_or(LoadingError::BadPath)
}
<commit_msg>Implement lazy load for themes<commit_after>// Copyright 2017 The xi-editor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Management of styles.

use std::collections::{HashMap, HashSet};
use std::ffi::OsStr;
use std::fs;
use std::iter::FromIterator;
use std::path::{Path, PathBuf};

use serde_json::{self, Value};
use syntect::dumps::{dump_to_file, from_dump_file};
use syntect::highlighting::StyleModifier as SynStyleModifier;
use syntect::highlighting::{Color, Highlighter, Theme, ThemeSet};
use syntect::LoadingError;

pub use syntect::highlighting::ThemeSettings;

pub const N_RESERVED_STYLES: usize = 8;
const SYNTAX_PRIORITY_DEFAULT: u16 = 200;
const SYNTAX_PRIORITY_LOWEST: u16 = 0;
pub const DEFAULT_THEME: &str = "InspiredGitHub";

#[derive(Clone, PartialEq, Eq, Default, Hash, Debug, Serialize, Deserialize)]
/// A mergeable style. All values except priority are optional.
///
/// Note: A `None` value represents the absense of preference; in the case of
/// boolean options, `Some(false)` means that this style will override a lower
/// priority value in the same field.
pub struct Style {
    /// The priority of this style, in the range (0, 1000). Used to resolve
    /// conflicting fields when merging styles. The higher priority wins.
    #[serde(skip_serializing)]
    pub priority: u16,
    /// The foreground text color, in ARGB.
    pub fg_color: Option<u32>,
    /// The background text color, in ARGB.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bg_color: Option<u32>,
    /// The font-weight, in the range 100-900, interpreted like the CSS
    /// font-weight property.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub weight: Option<u16>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub underline: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub italic: Option<bool>,
}

impl Style {
    /// Creates a new `Style` by converting from a `Syntect::StyleModifier`.
    pub fn from_syntect_style_mod(style: &SynStyleModifier) -> Self {
        let font_style = style.font_style.map(|s| s.bits()).unwrap_or_default();
        let weight = if (font_style & 1) != 0 { Some(700) } else { None };
        let underline = if (font_style & 2) != 0 { Some(true) } else { None };
        let italic = if (font_style & 4) != 0 { Some(true) } else { None };

        Self::new(
            SYNTAX_PRIORITY_DEFAULT,
            style.foreground.map(Self::rgba_from_syntect_color),
            style.background.map(Self::rgba_from_syntect_color),
            weight,
            underline,
            italic,
        )
    }

    pub fn new<O32, O16, OB>(
        priority: u16,
        fg_color: O32,
        bg_color: O32,
        weight: O16,
        underline: OB,
        italic: OB,
    ) -> Self
    where
        O32: Into<Option<u32>>,
        O16: Into<Option<u16>>,
        OB: Into<Option<bool>>,
    {
        assert!(priority <= 1000);
        Style {
            priority,
            fg_color: fg_color.into(),
            bg_color: bg_color.into(),
            weight: weight.into(),
            underline: underline.into(),
            italic: italic.into(),
        }
    }

    /// Returns the default style for the given `Theme`.
    pub fn default_for_theme(theme: &Theme) -> Self {
        let fg = theme.settings.foreground.unwrap_or(Color::BLACK);
        Style::new(
            SYNTAX_PRIORITY_LOWEST,
            Some(Self::rgba_from_syntect_color(fg)),
            None,
            None,
            None,
            None,
        )
    }

    /// Creates a new style by combining attributes of `self` and `other`.
    /// If both styles define an attribute, the highest priority wins; `other`
    /// wins in the case of a tie.
    ///
    /// Note: when merging multiple styles, apply them in increasing priority.
    pub fn merge(&self, other: &Style) -> Style {
        let (p1, p2) = if self.priority > other.priority { (self, other) } else { (other, self) };

        Style::new(
            p1.priority,
            p1.fg_color.or(p2.fg_color),
            p1.bg_color.or(p2.bg_color),
            p1.weight.or(p2.weight),
            p1.underline.or(p2.underline),
            p1.italic.or(p2.italic),
        )
    }

    /// Encode this `Style`, setting the `id` property.
    ///
    /// Note: this should only be used when sending the `def_style` RPC.
    pub fn to_json(&self, id: usize) -> Value {
        let mut as_val = serde_json::to_value(self).expect("failed to encode style");
        as_val["id"] = id.into();
        as_val
    }

    fn rgba_from_syntect_color(color: Color) -> u32 {
        let Color { r, g, b, a } = color;
        ((a as u32) << 24) | ((r as u32) << 16) | ((g as u32) << 8) | (b as u32)
    }
}

/// A map from styles to client identifiers for a given `Theme`.
pub struct ThemeStyleMap {
    themes: ThemeSet,
    theme_name: String,
    theme: Theme,
    default_style: Style,
    map: HashMap<Style, usize>,
    path_map: HashMap<String, PathBuf>,

    // It's not obvious we actually have to store the style, we seem to only need it
    // as the key in the map.
    styles: Vec<Style>,
    themes_dir: Option<PathBuf>,
    cache_dir: Option<PathBuf>,
    caching_enabled: bool,

    // Maintaining all theme paths for comparison on an fs event.
    state: HashSet<PathBuf>,
}

impl ThemeStyleMap {
    pub fn new(themes_dir: Option<PathBuf>) -> ThemeStyleMap {
        let themes = ThemeSet::load_defaults();
        let theme_name = DEFAULT_THEME.to_owned();
        let theme = themes.themes.get(&theme_name).expect("missing theme").to_owned();
        let default_style = Style::default_for_theme(&theme);
        let cache_dir = None;
        let caching_enabled = true;

        ThemeStyleMap {
            themes,
            theme_name,
            theme,
            default_style,
            map: HashMap::new(),
            path_map: HashMap::new(),
            styles: Vec::new(),
            themes_dir,
            cache_dir,
            caching_enabled,
            state: HashSet::new(),
        }
    }

    pub fn get_default_style(&self) -> &Style {
        &self.default_style
    }

    pub fn get_highlighter(&self) -> Highlighter {
        Highlighter::new(&self.theme)
    }

    pub fn get_theme_name(&self) -> &str {
        &self.theme_name
    }

    pub fn get_theme_settings(&self) -> &ThemeSettings {
        &self.theme.settings
    }

    pub fn get_theme_names(&self) -> Vec<String> {
        self.themes.themes.keys().cloned().collect()
    }

    pub fn contains_theme(&self, k: &str) -> bool {
        self.themes.themes.contains_key(k)
    }

    pub fn set_theme(&mut self, theme_name: &str) -> Result<(), &'static str> {
        // If we haven't loaded the theme before, we try to load it from the dump (if the dump
        // exists), or load it from the file itself.
        // Otherwise, we just load the cached theme from our theme map.
        if !self.contains_theme(theme_name) {
            let theme_p = &self.path_map.get(theme_name).cloned();
            if let Some(theme_p) = theme_p {
                match self.try_load_from_dump(theme_p) {
                    // If loading from the dump is successful, this will get us the theme name
                    // and its theme data.
                    Some((dump_theme_name, dump_theme_data)) => {
                        self.insert_to_map(dump_theme_name, dump_theme_data, theme_p);
                    }
                    None => {
                        let _ = self.load_theme(theme_p);
                    }
                }
            } else {
                return Err("can't find path for theme");
            }
        }

        if let Some(new_theme) = self.themes.themes.get(theme_name) {
            self.theme = new_theme.to_owned();
            self.theme_name = theme_name.to_owned();
            self.default_style = Style::default_for_theme(&self.theme);
            self.map = HashMap::new();
            self.styles = Vec::new();
            Ok(())
        } else {
            Err("unknown theme")
        }
    }

    pub fn merge_with_default(&self, style: &Style) -> Style {
        self.default_style.merge(style)
    }

    pub fn lookup(&self, style: &Style) -> Option<usize> {
        self.map.get(style).cloned()
    }

    pub fn add(&mut self, style: &Style) -> usize {
        let result = self.styles.len() + N_RESERVED_STYLES;
        self.map.insert(style.clone(), result);
        self.styles.push(style.clone());
        result
    }

    /// Delete key and the corresponding dump file from the themes map.
    pub(crate) fn remove_theme(&mut self, path: &Path) -> Option<String> {
        validate_theme_file(path).ok()?;

        let theme_name = path.file_stem().and_then(OsStr::to_str)?;
        self.themes.themes.remove(theme_name);
        self.state.remove(path);

        let dump_p = self.get_dump_path(theme_name)?;
        if dump_p.exists() {
            let _ = fs::remove_file(dump_p);
        }

        Some(theme_name.to_string())
    }

    /// Cache all themes names and their paths inside the given directory.
    pub(crate) fn load_theme_dir(&mut self) {
        if let Some(themes_dir) = self.themes_dir.clone() {
            match ThemeSet::discover_theme_paths(themes_dir) {
                Ok(themes) => {
                    self.caching_enabled = self.caching_enabled && self.init_cache_dir();
<<<<<<< HEAD

                    for theme_p in &themes {
                        match self.try_load_from_dump(theme_p) {
                            Some((k, v)) => {
                                self.insert_to_map(k, v, theme_p);
                            }
                            None => {
                                let _ = self.load_theme(theme_p);
                            }
=======
                    // We look through the theme folder here and cache their names/paths to a
                    // path hashmap.
                    for theme_p in themes.iter() {
                        if let Some(theme_name) = theme_p.file_stem().and_then(OsStr::to_str) {
                            self.path_map.insert(theme_name.to_string(), theme_p.to_path_buf());
                        } else {
                            error!("Invalid theme name, skipping");
>>>>>>> Implement lazy load for themes
                        }
                    }
                }
                Err(e) => error!("Error loading themes dir: {:?}", e),
            }
        }
    }

    /// A wrapper around `from_dump_file`
    /// to validate the state of dump file.
    /// Invalidates if mod time of dump is less
    /// than the original one.
    fn try_load_from_dump(&self, theme_p: &Path) -> Option<(String, Theme)> {
        if !self.caching_enabled {
            return None;
        }

        let theme_name = theme_p.file_stem().and_then(OsStr::to_str)?;

        let dump_p = self.get_dump_path(theme_name)?;

        if !&dump_p.exists() {
            return None;
        }

        //NOTE: `try_load_from_dump` will return `None` if the file at
        //`dump_p` or `theme_p` is deleted before the execution of this fn.
        let mod_t = fs::metadata(&dump_p).and_then(|md| md.modified()).ok()?;
        let mod_t_orig = fs::metadata(theme_p).and_then(|md| md.modified()).ok()?;

        if mod_t >= mod_t_orig {
            from_dump_file(&dump_p).ok().map(|t| (theme_name.to_owned(), t))
        } else {
            // Delete dump file
            let _ = fs::remove_file(&dump_p);
            None
        }
    }

    /// Loads theme using syntect's `get_theme` fn to our `theme` map.
    /// Stores binary dump in a file with `tmdump` extension, only if
    /// caching is enabled.
    pub(crate) fn load_theme(&mut self, theme_p: &Path) -> Result<String, LoadingError> {
        validate_theme_file(theme_p)?;
        let theme = ThemeSet::get_theme(theme_p)?;
        let theme_name =
            theme_p.file_stem().and_then(OsStr::to_str).ok_or(LoadingError::BadPath)?;

        if self.caching_enabled {
            if let Some(dump_p) = self.get_dump_path(theme_name) {
                let _ = dump_to_file(&theme, dump_p);
            }
        }
        self.insert_to_map(theme_name.to_owned(), theme, theme_p);
        Ok(theme_name.to_owned())
    }

    fn insert_to_map(&mut self, k: String, v: Theme, p: &Path) {
        self.themes.themes.insert(k, v);

        //Maintain a record for future syncing
        self.state.insert(p.to_path_buf());
    }

    /// Returns dump's path corresponding to the given theme name.
    fn get_dump_path(&self, theme_name: &str) -> Option<PathBuf> {
        self.cache_dir.as_ref().map(|p| p.join(theme_name).with_extension("tmdump"))
    }

    /// Compare the stored file paths in `self.state`
    /// to the present ones.
    pub(crate) fn sync_dir(&mut self, dir: Option<&Path>) {
        if let Some(themes_dir) = dir {
            if let Ok(paths) = ThemeSet::discover_theme_paths(themes_dir) {
                let current_state = HashSet::from_iter(paths.into_iter());
                let maintained_state = self.state.clone();

                let to_insert = current_state.difference(&maintained_state);
                for path in to_insert {
                    let _ = self.load_theme(path);
                }

                let to_remove = maintained_state.difference(&current_state);
                for path in to_remove {
                    self.remove_theme(path);
                }
            }
        }
    }

    /// Creates the cache dir returns true
    /// if it is successfully initialized or
    /// already exists.
    fn init_cache_dir(&mut self) -> bool {
        self.cache_dir = self.themes_dir.clone().map(|p| p.join("cache"));

        if let Some(ref p) = self.cache_dir {
            if p.exists() {
                return true;
            }
            fs::DirBuilder::new().create(&p).is_ok()
        } else {
            false
        }
    }
}

/// Used to remove files with extension other than `tmTheme`.
fn validate_theme_file(path: &Path) -> Result<bool, LoadingError> {
    path.extension().map(|e| e != "tmTheme").ok_or(LoadingError::BadPath)
}
<|endoftext|>


<commit_before>use std::{collections::BTreeMap, panic};

use quickcheck::{Arbitrary, Gen, RngCore};
use rand::distributions::{Distribution, Gamma};
use rand::{Rng, SeedableRng, StdRng};

use sled::*;

#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq)]
pub struct Key(pub Vec<u8>);

struct SledGen {
    r: StdRng,
    size: usize,
}

impl Gen for SledGen {
    fn size(&self) -> usize {
        self.size
    }
}

impl RngCore for SledGen {
    fn next_u32(&mut self) -> u32 {
        self.r.gen::<u32>()
    }

    fn next_u64(&mut self) -> u64 {
        self.r.gen::<u64>()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.r.fill_bytes(dest)
    }

    fn try_fill_bytes(
        &mut self,
        dest: &mut [u8],
    ) -> std::result::Result<(), rand::Error> {
        self.r.try_fill_bytes(dest)
    }
}

pub fn fuzz_then_shrink(buf: &[u8]) {
    let ops: Vec<Op> = buf
        .chunks(2)
        .map(|chunk| {
            let mut seed = [0u8; 32];
            seed[0..chunk.len()].copy_from_slice(chunk);
            let rng: StdRng = SeedableRng::from_seed(seed);
            let mut sled_rng = SledGen { r: rng, size: 2 };

            Op::arbitrary(&mut sled_rng)
        }).collect();

    println!("working on ops {:?}", ops);
    match panic::catch_unwind(move || {
        prop_tree_matches_btreemap(ops, 0, 100, false)
    }) {
        Ok(_) => {}
        Err(_e) => panic!("TODO"),
    }
}

impl Arbitrary for Key {
    fn arbitrary<G: Gen>(g: &mut G) -> Key {
        let gs = g.size();
        let gamma = Gamma::new(0.3, gs as f64);
        let v = gamma.sample(&mut rand::thread_rng());
        let len = if v > 30000.0 {
            30000
        } else if v < 1. && v > 0.0001 {
            1
        } else {
            v as usize
        };

        let space = g.gen_range(0, gs) + 1;

        let inner =
            (0..len).map(|_| g.gen_range(0, space) as u8).collect();

        Key(inner)
    }

    fn shrink(&self) -> Box<Iterator<Item = Key>> {
        let mut out = vec![];
        for len in 0..self.0.len() {
            out.push(Key(self.0[..len]
                .iter()
                .map(|e| e / 2)
                .collect()));
            out.push(Key(self.0[..len].to_vec()));
        }
        Box::new(out.into_iter())
    }
}

#[derive(Debug, Clone)]
pub enum Op {
    Set(Key, u8),
    Merge(Key, u8),
    Get(Key),
    Del(Key),
    Cas(Key, u8, u8),
    Scan(Key, usize),
    Restart,
}

use self::Op::*;

impl Arbitrary for Op {
    fn arbitrary<G: Gen>(g: &mut G) -> Op {
        if g.gen_bool(1. / 10.) {
            return Restart;
        }

        let choice = g.gen_range(0, 6);

        match choice {
            0 => Set(Key::arbitrary(g), g.gen::<u8>()),
            1 => Merge(Key::arbitrary(g), g.gen::<u8>()),
            2 => Get(Key::arbitrary(g)),
            3 => Del(Key::arbitrary(g)),
            4 => Cas(Key::arbitrary(g), g.gen::<u8>(), g.gen::<u8>()),
            5 => Scan(Key::arbitrary(g), g.gen_range::<usize>(0, 40)),
            _ => panic!("impossible choice"),
        }
    }

    fn shrink(&self) -> Box<Iterator<Item = Op>> {
        match *self {
            Set(ref k, v) => {
                Box::new(k.shrink().map(move |sk| Set(sk, v.clone())))
            }
            Merge(ref k, v) => {
                Box::new(k.shrink().map(move |k| Merge(k, v.clone())))
            }
            Get(ref k) => Box::new(k.shrink().map(move |k| Get(k))),
            Cas(ref k, old, new) => Box::new(Box::new(
                k.shrink()
                    .map(move |k| Cas(k, old.clone(), new.clone())),
            )),
            Scan(ref k, len) => {
                Box::new(k.shrink().map(move |k| Scan(k, len)))
            }
            Del(ref k) => Box::new(k.shrink().map(|k| Del(k))),
            Restart => Box::new(vec![].into_iter()),
        }
    }
}

fn bytes_to_u16(v: &[u8]) -> u16 {
    assert_eq!(v.len(), 2);
    ((v[0] as u16) << 8) + v[1] as u16
}

fn u16_to_bytes(u: u16) -> Vec<u8> {
    let ret = vec![(u >> 8) as u8, u as u8];
    ret
}

// just adds up values as if they were u16's
fn test_merge_operator(
    _k: &[u8],
    old: Option<&[u8]>,
    to_merge: &[u8],
) -> Option<Vec<u8>> {
    let base = old.unwrap_or(&[0, 0]);
    let base_n = bytes_to_u16(base);
    let new_n = base_n + to_merge[0] as u16;
    let ret = u16_to_bytes(new_n);
    Some(ret)
}

pub fn prop_tree_matches_btreemap(
    ops: Vec<Op>,
    blink_node_exponent: u8,
    snapshot_after: u8,
    flusher: bool,
) -> bool {
    use self::*;
    let config = ConfigBuilder::new()
        .temporary(true)
        .snapshot_after_ops(snapshot_after as usize + 1)
        .flush_every_ms(if flusher { Some(1) } else { None })
        .io_buf_size(10000)
        .blink_node_split_size(
            1 << std::cmp::min(blink_node_exponent, 20),
        ).cache_capacity(40)
        .cache_bits(0)
        .merge_operator(test_merge_operator)
        .build();

    let mut tree = sled::Tree::start(config.clone()).unwrap();
    let mut reference: BTreeMap<Key, u16> = BTreeMap::new();

    for op in ops.into_iter() {
        match op {
            Set(k, v) => {
                tree.set(k.0.clone(), vec![0, v]).unwrap();
                reference.insert(k.clone(), v as u16);
            }
            Merge(k, v) => {
                tree.merge(k.0.clone(), vec![v]).unwrap();
                let mut entry = reference.entry(k).or_insert(0u16);
                *entry += v as u16;
            }
            Get(k) => {
                let res1 = tree
                    .get(&*k.0)
                    .unwrap()
                    .map(|v| bytes_to_u16(&*v));
                let res2 = reference.get(&k).cloned();
                assert_eq!(res1, res2);
            }
            Del(k) => {
                tree.del(&*k.0).unwrap();
                reference.remove(&k);
            }
            Cas(k, old, new) => {
                let tree_old = tree.get(&*k.0).unwrap();
                if let Some(old_tree) = tree_old {
                    if old_tree == &*vec![0, old] {
                        tree.set(k.0.clone(), vec![0, new]).unwrap();
                    }
                }

                let ref_old = reference.get(&k).cloned();
                if ref_old == Some(old as u16) {
                    reference.insert(k, new as u16);
                }
            }
            Scan(k, len) => {
                let mut tree_iter = tree
                    .scan(&*k.0)
                    .take(len)
                    .map(|res| res.unwrap());
                let ref_iter = reference
                    .iter()
                    .filter(|&(ref rk, _rv)| **rk >= k)
                    .take(len)
                    .map(|(ref rk, ref rv)| (rk.0.clone(), **rv));

                for r in ref_iter {
                    let tree_next = tree_iter.next().unwrap();
                    let lhs = (tree_next.0, &*tree_next.1);
                    let rhs = (r.0.clone(), &*u16_to_bytes(r.1));
                    assert_eq!(
                        lhs, rhs,
                        "expected iteration over the Tree \
                         to match our BTreeMap model"
                    );
                }
            }
            Restart => {
                drop(tree);
                tree = sled::Tree::start(config.clone()).unwrap();
            }
        }
    }

    true
}

#[test]
fn test_fuzz_test() {
    let seed = [0; 32];
    fuzz_then_shrink(&seed);
    let seed = [0; 31];
    fuzz_then_shrink(&seed);
    let seed = [0; 33];
    fuzz_then_shrink(&seed);
    fuzz_then_shrink(&[]);
}
<commit_msg>Fix warnings with new version of rand<commit_after>use std::{collections::BTreeMap, panic};

use quickcheck::{Arbitrary, Gen, RngCore};
use rand::distributions::{Distribution, Gamma};
use rand::{rngs::StdRng, Rng, SeedableRng};

use sled::*;

#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq)]
pub struct Key(pub Vec<u8>);

struct SledGen {
    r: StdRng,
    size: usize,
}

impl Gen for SledGen {
    fn size(&self) -> usize {
        self.size
    }
}

impl RngCore for SledGen {
    fn next_u32(&mut self) -> u32 {
        self.r.gen::<u32>()
    }

    fn next_u64(&mut self) -> u64 {
        self.r.gen::<u64>()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.r.fill_bytes(dest)
    }

    fn try_fill_bytes(
        &mut self,
        dest: &mut [u8],
    ) -> std::result::Result<(), rand::Error> {
        self.r.try_fill_bytes(dest)
    }
}

pub fn fuzz_then_shrink(buf: &[u8]) {
    let ops: Vec<Op> = buf
        .chunks(2)
        .map(|chunk| {
            let mut seed = [0u8; 32];
            seed[0..chunk.len()].copy_from_slice(chunk);
            let rng: StdRng = SeedableRng::from_seed(seed);
            let mut sled_rng = SledGen { r: rng, size: 2 };

            Op::arbitrary(&mut sled_rng)
        }).collect();

    println!("working on ops {:?}", ops);
    match panic::catch_unwind(move || {
        prop_tree_matches_btreemap(ops, 0, 100, false)
    }) {
        Ok(_) => {}
        Err(_e) => panic!("TODO"),
    }
}

impl Arbitrary for Key {
    fn arbitrary<G: Gen>(g: &mut G) -> Key {
        let gs = g.size();
        let gamma = Gamma::new(0.3, gs as f64);
        let v = gamma.sample(&mut rand::thread_rng());
        let len = if v > 30000.0 {
            30000
        } else if v < 1. && v > 0.0001 {
            1
        } else {
            v as usize
        };

        let space = g.gen_range(0, gs) + 1;

        let inner =
            (0..len).map(|_| g.gen_range(0, space) as u8).collect();

        Key(inner)
    }

    fn shrink(&self) -> Box<Iterator<Item = Key>> {
        let mut out = vec![];
        for len in 0..self.0.len() {
            out.push(Key(self.0[..len]
                .iter()
                .map(|e| e / 2)
                .collect()));
            out.push(Key(self.0[..len].to_vec()));
        }
        Box::new(out.into_iter())
    }
}

#[derive(Debug, Clone)]
pub enum Op {
    Set(Key, u8),
    Merge(Key, u8),
    Get(Key),
    Del(Key),
    Cas(Key, u8, u8),
    Scan(Key, usize),
    Restart,
}

use self::Op::*;

impl Arbitrary for Op {
    fn arbitrary<G: Gen>(g: &mut G) -> Op {
        if g.gen_bool(1. / 10.) {
            return Restart;
        }

        let choice = g.gen_range(0, 6);

        match choice {
            0 => Set(Key::arbitrary(g), g.gen::<u8>()),
            1 => Merge(Key::arbitrary(g), g.gen::<u8>()),
            2 => Get(Key::arbitrary(g)),
            3 => Del(Key::arbitrary(g)),
            4 => Cas(Key::arbitrary(g), g.gen::<u8>(), g.gen::<u8>()),
            5 => Scan(Key::arbitrary(g), g.gen_range(0, 40)),
            _ => panic!("impossible choice"),
        }
    }

    fn shrink(&self) -> Box<Iterator<Item = Op>> {
        match *self {
            Set(ref k, v) => {
                Box::new(k.shrink().map(move |sk| Set(sk, v.clone())))
            }
            Merge(ref k, v) => {
                Box::new(k.shrink().map(move |k| Merge(k, v.clone())))
            }
            Get(ref k) => Box::new(k.shrink().map(move |k| Get(k))),
            Cas(ref k, old, new) => Box::new(Box::new(
                k.shrink()
                    .map(move |k| Cas(k, old.clone(), new.clone())),
            )),
            Scan(ref k, len) => {
                Box::new(k.shrink().map(move |k| Scan(k, len)))
            }
            Del(ref k) => Box::new(k.shrink().map(|k| Del(k))),
            Restart => Box::new(vec![].into_iter()),
        }
    }
}

fn bytes_to_u16(v: &[u8]) -> u16 {
    assert_eq!(v.len(), 2);
    ((v[0] as u16) << 8) + v[1] as u16
}

fn u16_to_bytes(u: u16) -> Vec<u8> {
    let ret = vec![(u >> 8) as u8, u as u8];
    ret
}

// just adds up values as if they were u16's
fn test_merge_operator(
    _k: &[u8],
    old: Option<&[u8]>,
    to_merge: &[u8],
) -> Option<Vec<u8>> {
    let base = old.unwrap_or(&[0, 0]);
    let base_n = bytes_to_u16(base);
    let new_n = base_n + to_merge[0] as u16;
    let ret = u16_to_bytes(new_n);
    Some(ret)
}

pub fn prop_tree_matches_btreemap(
    ops: Vec<Op>,
    blink_node_exponent: u8,
    snapshot_after: u8,
    flusher: bool,
) -> bool {
    use self::*;
    let config = ConfigBuilder::new()
        .temporary(true)
        .snapshot_after_ops(snapshot_after as usize + 1)
        .flush_every_ms(if flusher { Some(1) } else { None })
        .io_buf_size(10000)
        .blink_node_split_size(
            1 << std::cmp::min(blink_node_exponent, 20),
        ).cache_capacity(40)
        .cache_bits(0)
        .merge_operator(test_merge_operator)
        .build();

    let mut tree = sled::Tree::start(config.clone()).unwrap();
    let mut reference: BTreeMap<Key, u16> = BTreeMap::new();

    for op in ops.into_iter() {
        match op {
            Set(k, v) => {
                tree.set(k.0.clone(), vec![0, v]).unwrap();
                reference.insert(k.clone(), v as u16);
            }
            Merge(k, v) => {
                tree.merge(k.0.clone(), vec![v]).unwrap();
                let mut entry = reference.entry(k).or_insert(0u16);
                *entry += v as u16;
            }
            Get(k) => {
                let res1 = tree
                    .get(&*k.0)
                    .unwrap()
                    .map(|v| bytes_to_u16(&*v));
                let res2 = reference.get(&k).cloned();
                assert_eq!(res1, res2);
            }
            Del(k) => {
                tree.del(&*k.0).unwrap();
                reference.remove(&k);
            }
            Cas(k, old, new) => {
                let tree_old = tree.get(&*k.0).unwrap();
                if let Some(old_tree) = tree_old {
                    if old_tree == &*vec![0, old] {
                        tree.set(k.0.clone(), vec![0, new]).unwrap();
                    }
                }

                let ref_old = reference.get(&k).cloned();
                if ref_old == Some(old as u16) {
                    reference.insert(k, new as u16);
                }
            }
            Scan(k, len) => {
                let mut tree_iter = tree
                    .scan(&*k.0)
                    .take(len)
                    .map(|res| res.unwrap());
                let ref_iter = reference
                    .iter()
                    .filter(|&(ref rk, _rv)| **rk >= k)
                    .take(len)
                    .map(|(ref rk, ref rv)| (rk.0.clone(), **rv));

                for r in ref_iter {
                    let tree_next = tree_iter.next().unwrap();
                    let lhs = (tree_next.0, &*tree_next.1);
                    let rhs = (r.0.clone(), &*u16_to_bytes(r.1));
                    assert_eq!(
                        lhs, rhs,
                        "expected iteration over the Tree \
                         to match our BTreeMap model"
                    );
                }
            }
            Restart => {
                drop(tree);
                tree = sled::Tree::start(config.clone()).unwrap();
            }
        }
    }

    true
}

#[test]
fn test_fuzz_test() {
    let seed = [0; 32];
    fuzz_then_shrink(&seed);
    let seed = [0; 31];
    fuzz_then_shrink(&seed);
    let seed = [0; 33];
    fuzz_then_shrink(&seed);
    fuzz_then_shrink(&[]);
}
<|endoftext|>


<commit_before>use chrono::DateTime;
use chrono::offset::Utc;
use diesel;
use diesel::prelude::*;
use diesel::pg::PgConnection;
use diesel::pg::upsert::*;
use futures::{Future, stream, Stream};
use futures_cpupool::CpuPool;
use hyper::{Body, Client, Uri};
use hyper::client::HttpConnector;
use hyper_tls::HttpsConnector;
use InsertError;
use kuchiki;
use kuchiki::traits::*;
use r2d2_diesel::ConnectionManager;
use r2d2::Pool;
use rusoto_core::{default_tls_client, Region};
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{PutObjectRequest, S3Client, S3};
use schema::ads;
use server::AdPost;

const ENDPOINT: &'static str = "https://pp-facebook-ads.s3.amazonaws.com/";

fn get_title(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    document
        .select("h5 a, h6 a, strong")
        .map_err(InsertError::HTML)?
        .nth(0)
        .and_then(|a| Some(a.text_contents()))
        .ok_or(InsertError::HTML(()))
}

fn get_image(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    document
        .select("img")
        .map_err(InsertError::HTML)?
        .nth(0)
        .and_then(|a| {
            a.attributes.borrow().get("src").and_then(
                |src| Some(src.to_string()),
            )
        })
        .ok_or(InsertError::HTML(()))
}

fn get_message(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    let selectors = vec![".userContent p", "span"];
    let iters = selectors.iter().map(|s| document.select(s)).flat_map(|a| a);

    iters
        .map(|i| {
            i.fold(String::new(), |m, a| m + &a.as_node().to_string())
        })
        .filter(|i| !i.is_empty())
        .nth(0)
        .ok_or(InsertError::HTML(()))
}

fn get_images(document: &kuchiki::NodeRef) -> Result<Vec<String>, InsertError> {
    Ok(
        document
            .select("img")
            .map_err(InsertError::HTML)?
            .skip(1)
            .map(|a| {
                a.attributes.borrow().get("src").and_then(
                    |s| Some(s.to_string()),
                )
            })
            .filter(|s| s.is_some())
            .map(|s| s.unwrap())
            .collect::<Vec<String>>(),
    )
}

#[derive(AsChangeset, Debug)]
#[table_name = "ads"]
pub struct Images {
    thumbnail: Option<String>,
    images: Vec<String>,
    title: String,
    message: String,
    html: String,
}

impl Images {
    fn from_ad(ad: &Ad, images: Vec<Uri>) -> Result<Images, InsertError> {
        let thumb = images
            .iter()
            .filter(|i| ad.thumbnail.contains(i.path()))
            .map(|i| ENDPOINT.to_string() + i.path().trim_left_matches('/'))
            .nth(0);

        let mut rest = images.clone();
        if let Some(thumb) = thumb.clone() {
            rest.retain(|x| !thumb.contains(x.path()))
        };

        let collection = rest.iter()
            .filter(|i| ad.images.iter().any(|a| a.contains(i.path())))
            .map(|i| ENDPOINT.to_string() + i.path().trim_left_matches('/'))
            .collect::<Vec<String>>();

        let document = kuchiki::parse_html().one(ad.html.clone());
        for a in document.select("img").map_err(InsertError::HTML)? {
            if let Some(x) = a.attributes.borrow_mut().get_mut("src") {
                if let Ok(u) = x.parse::<Uri>() {
                    if let Some(i) = images.iter().find(|i| i.path() == u.path()) {
                        *x = ENDPOINT.to_string() + i.path().trim_left_matches('/');
                    } else {
                        *x = "".to_string();
                    }
                }
            };
        }

        let title = get_title(&document)?;
        let message = get_title(&document)?;
        Ok(Images {
            thumbnail: thumb,
            images: collection,
            title: title,
            html: document
                .select("div")
                .map_err(InsertError::HTML)?
                .nth(0)
                .ok_or(InsertError::HTML(()))?
                .as_node()
                .to_string(),
            message: message,
        })
    }
}

#[derive(Queryable, Debug, Clone)]
pub struct Ad {
    id: String,
    html: String,
    political: i32,
    not_political: i32,
    fuzzy_id: Option<i32>,
    title: String,
    message: String,
    thumbnail: String,
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
    browser_lang: String,
    images: Vec<String>,
    impressions: i32,
}

impl Ad {
    // This will asynchronously save the images to s3 we may very well end up
    // dropping images, but I can't see any way around it right now. Also we
    // should think about splitting this up, but I'm fine -- if a little
    // embarassed about it -- right now. This function swallows errors, and
    // there's a chance we'll end up with no images at the end, but I think we
    // can handle that in the extension's UI.
    pub fn grab_and_store(
        &self,
        client: Client<HttpsConnector<HttpConnector>, Body>,
        db: &Pool<ConnectionManager<PgConnection>>,
        pool: CpuPool,
    ) -> Box<Future<Item = (), Error = ()>> {
        let ad = self.clone();
        let pool_s3 = pool.clone();
        let pool_db = pool.clone();
        let db = db.clone();
        let future = stream::iter(self.image_urls())
            // filter ones we already have in the db and ones we can verify as
            // coming from fb, we don't want to become a malware vector :)
            // currently we recache images we already have, but ok.
            .filter(|u| {
                info!("testing {:?}", u.host());
                match u.host() {
                    Some(h) => (h == "pp-facebook-ads.s3.amazonaws.com" || h.ends_with("fbcdn.net")),
                    None => false
                }
            })
            // grab image
            .and_then(move |img| {
                info!("getting {:?}", img);
                let cloned = img.clone();
                client
                    .get(img)
                    .and_then(|res| {
                        res.body().concat2().and_then(|chunk| Ok((chunk, cloned)))
                    })
                    .map_err(InsertError::Hyper)
            })
            // upload them to s3
            .and_then(move |tuple| {
                let pool = pool_s3.clone();
                // we do this in a worker thread because rusoto isn't on
                // Hyper async yet.
                pool.spawn_fn(move || {
                    let client =
                        S3Client::new(default_tls_client().map_err(InsertError::TLS)?,
                                      DefaultCredentialsProvider::new().map_err(InsertError::AWS)?,
                                      Region::UsEast1);
                    let req = PutObjectRequest {
                        bucket: "pp-facebook-ads".to_string(),
                        key: tuple.1.path().trim_left_matches('/').to_string(),
                        acl: Some("public-read".to_string()),
                        body: Some(tuple.0.to_vec()),
                        ..PutObjectRequest::default()
                    };
                    client.put_object(&req).map_err(InsertError::S3)?;
                    Ok(tuple.1)
                })
            })
            .collect()
            // save the new urls to the database. the images variable will
            // include only those that we've successfully saved to s3, so we
            // have to do a funky merge here.
            .and_then(move |images| {
                let imgs = images.clone();
                pool_db.spawn_fn(move || {
                    use schema::ads::dsl::*;
                    let update = Images::from_ad(&ad, imgs)?;
                    let connection = db.get().map_err(InsertError::Timeout)?;
                    diesel::update(ads.filter(id.eq(&ad.id)))
                        .set(&update)
                        .execute(&*connection)
                        .map_err(InsertError::DataBase)?;
                    info!("saved {:?}", ad.id);
                    Ok(())
                })
            })
            .map_err(|e| {
                warn!("{:?}", e);
                ()
            });
        Box::new(future)
    }

    pub fn image_urls(&self) -> Vec<Result<Uri, InsertError>> {
        let images = [vec![self.thumbnail.clone()], self.images.clone()];
        images
            .concat()
            .iter()
            .map(|a| {
                let a: Uri = a.parse().map_err(InsertError::Uri)?;
                Ok(a)
            })
            .collect()
    }
}

#[derive(Insertable)]
#[table_name = "ads"]
pub struct NewAd<'a> {
    id: &'a str,
    html: &'a str,
    political: i32,
    not_political: i32,

    title: String,
    message: String,
    thumbnail: String,

    browser_lang: &'a str,
    images: Vec<String>,
    impressions: i32,
}

impl<'a> NewAd<'a> {
    pub fn new(ad: &'a AdPost) -> Result<NewAd<'a>, InsertError> {
        let document = kuchiki::parse_html().one(ad.html.clone());

        let thumb = get_image(&document)?;
        let images = get_images(&document)?;
        let message = get_message(&document)?;
        let title = get_title(&document)?;

        Ok(NewAd {
            id: &ad.id,
            html: &ad.html,
            // we try unwrapping or we chose the false branch in both of these
            // cases to count impressions
            political: if ad.political.unwrap_or(false) { 1 } else { 0 },
            not_political: if !ad.political.unwrap_or(true) { 1 } else { 0 },
            title: title,
            message: message,
            thumbnail: thumb,
            browser_lang: &ad.browser_lang,
            images: images,
            impressions: if !ad.political.is_some() { 1 } else { 0 },
        })
    }

    pub fn save(&self, pool: &Pool<ConnectionManager<PgConnection>>) -> Result<Ad, InsertError> {
        use schema::ads;
        use schema::ads::dsl::*;
        let connection = pool.get().map_err(InsertError::Timeout)?;
        // increment impressions if this is a background save,
        // otherwise increment political counters
        let ad: Ad = diesel::insert(&self.on_conflict(
            id,
            do_update().set((
                political.eq(political + self.political),
                not_political.eq(
                    not_political +
                        self.not_political,
                ),
                impressions.eq(
                    impressions + self.impressions,
                ),
                updated_at.eq(Utc::now()),
            )),
        )).into(ads::table)
            .get_result(&*connection)
            .map_err(InsertError::DataBase)?;
        Ok(ad)
    }
}

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

    #[test]
    fn ad_parsing() {
        let ad = include_str!("./html-test.txt");
        let post = AdPost {
            id: "test".to_string(),
            html: ad.to_string(),
            political: None,
            browser_lang: "en-US".to_string(),
        };
        let new_ad = NewAd::new(&post).unwrap();
        assert!(new_ad.thumbnail.len() > 0);
        assert_eq!(new_ad.images.len(), 2);
        assert!(new_ad.title.len() > 0);

        assert_eq!(
            kuchiki::parse_html().one(new_ad.message).text_contents(),
            kuchiki::parse_html().one("
#Valerian is “the best experience since ‘Avatar.’” See it in 3D and RealD3D theaters this Friday. Get tickets now: ValerianTickets.com
").text_contents()
        );
    }

    #[test]
    fn image_parsing() {
        use chrono::prelude::*;
        let ad = include_str!("./html-test.txt");
        let document = kuchiki::parse_html().one(ad);

        let saved_ad = Ad {
            id: "test".to_string(),
            html: ad.to_string(),
            political: 1,
            not_political: 2,
            fuzzy_id: None,
            title: get_title(&document).unwrap(),
            message: get_message(&document).unwrap(),
            thumbnail: get_image(&document).unwrap(),
            created_at: Utc::now(),
            updated_at: Utc::now(),
            browser_lang: "en-US".to_string(),
            images: get_images(&document).unwrap(),
            impressions: 1,
        };
        let urls = saved_ad
            .image_urls()
            .into_iter()
            .map(|x| x.unwrap())
            .collect();
        let images = Images::from_ad(&saved_ad, urls).unwrap();
        assert!(images.html != saved_ad.html);
        assert!(!images.html.contains("fbcdn"));
        assert!(!images.html.contains("html"));
        assert!(images.images.len() == saved_ad.images.len());
    }
}
<commit_msg>don't reupload images we already have<commit_after>use chrono::DateTime;
use chrono::offset::Utc;
use diesel;
use diesel::prelude::*;
use diesel::pg::PgConnection;
use diesel::pg::upsert::*;
use futures::{Future, stream, Stream};
use futures_cpupool::CpuPool;
use hyper::{Body, Client, Uri};
use hyper::client::HttpConnector;
use hyper_tls::HttpsConnector;
use InsertError;
use kuchiki;
use kuchiki::traits::*;
use r2d2_diesel::ConnectionManager;
use r2d2::Pool;
use rusoto_core::{default_tls_client, Region};
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{PutObjectRequest, S3Client, S3};
use schema::ads;
use server::AdPost;

const ENDPOINT: &'static str = "https://pp-facebook-ads.s3.amazonaws.com/";

fn get_title(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    document
        .select("h5 a, h6 a, strong")
        .map_err(InsertError::HTML)?
        .nth(0)
        .and_then(|a| Some(a.text_contents()))
        .ok_or(InsertError::HTML(()))
}

fn get_image(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    document
        .select("img")
        .map_err(InsertError::HTML)?
        .nth(0)
        .and_then(|a| {
            a.attributes.borrow().get("src").and_then(
                |src| Some(src.to_string()),
            )
        })
        .ok_or(InsertError::HTML(()))
}

fn get_message(document: &kuchiki::NodeRef) -> Result<String, InsertError> {
    let selectors = vec![".userContent p", "span"];
    let iters = selectors.iter().map(|s| document.select(s)).flat_map(|a| a);

    iters
        .map(|i| {
            i.fold(String::new(), |m, a| m + &a.as_node().to_string())
        })
        .filter(|i| !i.is_empty())
        .nth(0)
        .ok_or(InsertError::HTML(()))
}

fn get_images(document: &kuchiki::NodeRef) -> Result<Vec<String>, InsertError> {
    Ok(
        document
            .select("img")
            .map_err(InsertError::HTML)?
            .skip(1)
            .map(|a| {
                a.attributes.borrow().get("src").and_then(
                    |s| Some(s.to_string()),
                )
            })
            .filter(|s| s.is_some())
            .map(|s| s.unwrap())
            .collect::<Vec<String>>(),
    )
}

#[derive(AsChangeset, Debug)]
#[table_name = "ads"]
pub struct Images {
    thumbnail: Option<String>,
    images: Vec<String>,
    title: String,
    message: String,
    html: String,
}

impl Images {
    fn from_ad(ad: &Ad, images: Vec<Uri>) -> Result<Images, InsertError> {
        let thumb = images
            .iter()
            .filter(|i| ad.thumbnail.contains(i.path()))
            .map(|i| ENDPOINT.to_string() + i.path().trim_left_matches('/'))
            .nth(0);

        let mut rest = images.clone();
        if let Some(thumb) = thumb.clone() {
            rest.retain(|x| !thumb.contains(x.path()))
        };

        let collection = rest.iter()
            .filter(|i| ad.images.iter().any(|a| a.contains(i.path())))
            .map(|i| ENDPOINT.to_string() + i.path().trim_left_matches('/'))
            .collect::<Vec<String>>();

        let document = kuchiki::parse_html().one(ad.html.clone());
        for a in document.select("img").map_err(InsertError::HTML)? {
            if let Some(x) = a.attributes.borrow_mut().get_mut("src") {
                if let Ok(u) = x.parse::<Uri>() {
                    if let Some(i) = images.iter().find(|i| i.path() == u.path()) {
                        *x = ENDPOINT.to_string() + i.path().trim_left_matches('/');
                    } else {
                        *x = "".to_string();
                    }
                }
            };
        }

        let title = get_title(&document)?;
        let message = get_title(&document)?;
        Ok(Images {
            thumbnail: thumb,
            images: collection,
            title: title,
            html: document
                .select("div")
                .map_err(InsertError::HTML)?
                .nth(0)
                .ok_or(InsertError::HTML(()))?
                .as_node()
                .to_string(),
            message: message,
        })
    }
}

#[derive(Queryable, Debug, Clone)]
pub struct Ad {
    id: String,
    html: String,
    political: i32,
    not_political: i32,
    fuzzy_id: Option<i32>,
    title: String,
    message: String,
    thumbnail: String,
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
    browser_lang: String,
    images: Vec<String>,
    impressions: i32,
}

impl Ad {
    // This will asynchronously save the images to s3 we may very well end up
    // dropping images, but I can't see any way around it right now. Also we
    // should think about splitting this up, but I'm fine -- if a little
    // embarassed about it -- right now. This function swallows errors, and
    // there's a chance we'll end up with no images at the end, but I think we
    // can handle that in the extension's UI.
    pub fn grab_and_store(
        &self,
        client: Client<HttpsConnector<HttpConnector>, Body>,
        db: &Pool<ConnectionManager<PgConnection>>,
        pool: CpuPool,
    ) -> Box<Future<Item = (), Error = ()>> {
        let ad = self.clone();
        let pool_s3 = pool.clone();
        let pool_db = pool.clone();
        let db = db.clone();
        let future = stream::iter(self.image_urls())
            // filter ones we already have in the db and ones we can verify as
            // coming from fb, we don't want to become a malware vector :)
            // currently we redownload images we already have, but ok.
            .filter(|u| {
                info!("testing {:?}", u.host());
                match u.host() {
                    Some(h) => (h == "pp-facebook-ads.s3.amazonaws.com" || h.ends_with("fbcdn.net")),
                    None => false
                }
            })
            // grab image
            .and_then(move |img| {
                info!("getting {:?}", img);
                let cloned = img.clone();
                client
                    .get(img)
                    .and_then(|res| {
                        res.body().concat2().and_then(|chunk| Ok((chunk, cloned)))
                    })
                    .map_err(InsertError::Hyper)
            })
            // upload them to s3
            .and_then(move |tuple| {
                let pool = pool_s3.clone();
                // we do this in a worker thread because rusoto isn't on
                // Hyper async yet.
                pool.spawn_fn(move || {
                    if tuple.1.host().unwrap() != "pp-facebook-ads.s3.amazonaws.com" {
                        let client =
                            S3Client::new(default_tls_client().map_err(InsertError::TLS)?,
                                          DefaultCredentialsProvider::new().map_err(InsertError::AWS)?,
                                          Region::UsEast1);
                        let req = PutObjectRequest {
                            bucket: "pp-facebook-ads".to_string(),
                            key: tuple.1.path().trim_left_matches('/').to_string(),
                            acl: Some("public-read".to_string()),
                            body: Some(tuple.0.to_vec()),
                            ..PutObjectRequest::default()
                        };
                        client.put_object(&req).map_err(InsertError::S3)?;
                    }
                    Ok(tuple.1)
                })
            })
            .collect()
            // save the new urls to the database. the images variable will
            // include only those that we've successfully saved to s3, so we
            // have to do a funky merge here.
            .and_then(move |images| {
                let imgs = images.clone();
                pool_db.spawn_fn(move || {
                    use schema::ads::dsl::*;
                    let update = Images::from_ad(&ad, imgs)?;
                    let connection = db.get().map_err(InsertError::Timeout)?;
                    diesel::update(ads.filter(id.eq(&ad.id)))
                        .set(&update)
                        .execute(&*connection)
                        .map_err(InsertError::DataBase)?;
                    info!("saved {:?}", ad.id);
                    Ok(())
                })
            })
            .map_err(|e| {
                warn!("{:?}", e);
                ()
            });
        Box::new(future)
    }

    pub fn image_urls(&self) -> Vec<Result<Uri, InsertError>> {
        let images = [vec![self.thumbnail.clone()], self.images.clone()];
        images
            .concat()
            .iter()
            .map(|a| {
                let a: Uri = a.parse().map_err(InsertError::Uri)?;
                Ok(a)
            })
            .collect()
    }
}

#[derive(Insertable)]
#[table_name = "ads"]
pub struct NewAd<'a> {
    id: &'a str,
    html: &'a str,
    political: i32,
    not_political: i32,

    title: String,
    message: String,
    thumbnail: String,

    browser_lang: &'a str,
    images: Vec<String>,
    impressions: i32,
}

impl<'a> NewAd<'a> {
    pub fn new(ad: &'a AdPost) -> Result<NewAd<'a>, InsertError> {
        let document = kuchiki::parse_html().one(ad.html.clone());

        let thumb = get_image(&document)?;
        let images = get_images(&document)?;
        let message = get_message(&document)?;
        let title = get_title(&document)?;

        Ok(NewAd {
            id: &ad.id,
            html: &ad.html,
            // we try unwrapping or we chose the false branch in both of these
            // cases to count impressions
            political: if ad.political.unwrap_or(false) { 1 } else { 0 },
            not_political: if !ad.political.unwrap_or(true) { 1 } else { 0 },
            title: title,
            message: message,
            thumbnail: thumb,
            browser_lang: &ad.browser_lang,
            images: images,
            impressions: if !ad.political.is_some() { 1 } else { 0 },
        })
    }

    pub fn save(&self, pool: &Pool<ConnectionManager<PgConnection>>) -> Result<Ad, InsertError> {
        use schema::ads;
        use schema::ads::dsl::*;
        let connection = pool.get().map_err(InsertError::Timeout)?;
        // increment impressions if this is a background save,
        // otherwise increment political counters
        let ad: Ad = diesel::insert(&self.on_conflict(
            id,
            do_update().set((
                political.eq(political + self.political),
                not_political.eq(
                    not_political +
                        self.not_political,
                ),
                impressions.eq(
                    impressions + self.impressions,
                ),
                updated_at.eq(Utc::now()),
            )),
        )).into(ads::table)
            .get_result(&*connection)
            .map_err(InsertError::DataBase)?;
        Ok(ad)
    }
}

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

    #[test]
    fn ad_parsing() {
        let ad = include_str!("./html-test.txt");
        let post = AdPost {
            id: "test".to_string(),
            html: ad.to_string(),
            political: None,
            browser_lang: "en-US".to_string(),
        };
        let new_ad = NewAd::new(&post).unwrap();
        assert!(new_ad.thumbnail.len() > 0);
        assert_eq!(new_ad.images.len(), 2);
        assert!(new_ad.title.len() > 0);

        assert_eq!(
            kuchiki::parse_html().one(new_ad.message).text_contents(),
            kuchiki::parse_html().one("
#Valerian is “the best experience since ‘Avatar.’” See it in 3D and RealD3D theaters this Friday. Get tickets now: ValerianTickets.com
").text_contents()
        );
    }

    #[test]
    fn image_parsing() {
        use chrono::prelude::*;
        let ad = include_str!("./html-test.txt");
        let document = kuchiki::parse_html().one(ad);

        let saved_ad = Ad {
            id: "test".to_string(),
            html: ad.to_string(),
            political: 1,
            not_political: 2,
            fuzzy_id: None,
            title: get_title(&document).unwrap(),
            message: get_message(&document).unwrap(),
            thumbnail: get_image(&document).unwrap(),
            created_at: Utc::now(),
            updated_at: Utc::now(),
            browser_lang: "en-US".to_string(),
            images: get_images(&document).unwrap(),
            impressions: 1,
        };
        let urls = saved_ad
            .image_urls()
            .into_iter()
            .map(|x| x.unwrap())
            .collect();
        let images = Images::from_ad(&saved_ad, urls).unwrap();
        assert!(images.html != saved_ad.html);
        assert!(!images.html.contains("fbcdn"));
        assert!(!images.html.contains("html"));
        assert!(images.images.len() == saved_ad.images.len());
    }
}
<|endoftext|>


<commit_before>// need to raise recursion limit because many combinators
#![recursion_limit = "128"]
#![deny(bare_trait_objects)]

use std::io::{BufRead as _, BufReader, Write as _};
use std::path::PathBuf;

use lazy_static::lazy_static;
use regex::Regex;
use serde::{Deserialize, Serialize};
use serde_json;
use structopt::StructOpt;

use futures::{future, stream, sync::mpsc::channel, Future as _, Sink as _, Stream as _};
use log::{error, info, warn};

use gerritbot as bot;
use gerritbot_gerrit as gerrit;
use gerritbot_spark as spark;

#[derive(StructOpt, Debug)]
#[structopt(name = "gerritbot-console", rename_all = "kebab-case")]
/// Run the gerritbot without actually connecting to Spark. Instead the
/// program's stdin can be used to simulate sending messages. Replies will be
/// sent to stdout. Log messages will only appear on stderr. If the --email
/// option (see below) is given each input line will be treated as message from
/// that user. Otherwise each line has to be prefixed with the email address of
/// the intended sender, a colon (:) and an optional space character.
struct Args {
    /// Gerrit bot username
    ///
    /// This is the username the bot uses to connect to the Gerrit server.
    #[structopt(short, long)]
    username: String,
    /// Gerrit hostname
    ///
    /// Address of the Gerrit server.
    hostname: String,
    /// Gerrit SSH port
    #[structopt(short, long, default_value = "29418")]
    port: u32,
    /// Path to SSH private key
    #[structopt(short, long, parse(from_os_str))]
    identity_file: PathBuf,
    /// Enable verbose output
    #[structopt(short, long)]
    verbose: bool,
    /// Log only errors
    #[structopt(short, long, conflicts_with = "verbose")]
    quiet: bool,
    /// User email address
    ///
    /// If given input messages will be treated as if coming from this user.
    /// Additionally, an "enable" message will be injected before the first
    /// input line.
    ///
    /// This option is mutually exclusive with the --json option.
    #[structopt(short, long, conflicts_with = "json")]
    email: Option<String>,
    /// JSON mode
    ///
    /// Switch input and output to JSON. This is option mutually exclusive with
    /// the --email option.
    #[structopt(long)]
    json: bool,
    /// Lua formatting script
    ///
    /// Can be used to change or test the formatting of messages. If not present
    /// the internal default format script will be used.
    #[structopt(long = "format-script")]
    format_script: Option<String>,
}

#[derive(Deserialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
struct SimpleInputMessage {
    person_email: spark::Email,
    person_id: Option<spark::PersonId>,
    text: String,
}

fn email_to_person_id(email: &spark::Email) -> spark::PersonId {
    spark::PersonId::new(email.to_string())
}

impl Into<spark::Message> for SimpleInputMessage {
    fn into(self) -> spark::Message {
        let SimpleInputMessage {
            person_email,
            person_id,
            text,
        } = self;
        let person_id = person_id.unwrap_or_else(|| email_to_person_id(&person_email));
        spark::Message::test_message(person_email, person_id, text)
    }
}

#[derive(Serialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
struct SimpleOutputMessage {
    person_id: spark::PersonId,
    text: String,
}

#[derive(Clone)]
enum ConsoleSparkClient {
    Plain,
    Json,
}

impl bot::SparkClient for ConsoleSparkClient {
    type ReplyFuture = future::FutureResult<(), spark::Error>;
    fn reply(&self, person_id: &spark::PersonId, msg: &str) -> Self::ReplyFuture {
        // Write synchronously and crash if writing fails. There's no point in
        // error handling here.
        match self {
            ConsoleSparkClient::Plain => write!(std::io::stdout(), "{}: {}\n", person_id, msg)
                .expect("writing to stdout failed"),
            ConsoleSparkClient::Json => {
                let message = SimpleOutputMessage {
                    person_id: person_id.clone(),
                    text: msg.to_string(),
                };
                serde_json::to_writer(std::io::stdout(), &message)
                    .expect("writing JSON to stdout failed");
                std::io::stdout()
                    .write(b"\n")
                    .expect("writing to stdout failed");
            }
        }
        future::ok(())
    }
}

fn main() {
    let args = Args::from_args();
    stderrlog::new()
        .module(module_path!())
        .module("gerritbot")
        .module("gerritbot_gerrit")
        .timestamp(stderrlog::Timestamp::Second)
        .verbosity(match (args.quiet, args.verbose) {
            (true, _) => 0,      // ERROR
            (false, false) => 2, // INFO
            (_, true) => 4,      // TRACE
        })
        .init()
        .unwrap();

    let connect_to_gerrit = || {
        info!(
            "Connecting to gerrit with username {} at {}:{}",
            args.username, args.hostname, args.port,
        );
        gerrit::Connection::connect(
            format!("{}:{}", args.hostname, args.port),
            args.username.clone(),
            args.identity_file.clone(),
        )
        .unwrap_or_else(|e| {
            error!("failed to connect to gerrit: {}", e);
            std::process::exit(1);
        })
    };
    let gerrit_event_stream = gerrit::extended_event_stream(
        connect_to_gerrit(),
        connect_to_gerrit(),
        bot::request_extended_gerrit_info,
    );
    let gerrit_command_runner = gerrit::CommandRunner::new(connect_to_gerrit());
    let bot_builder = bot::Builder::new(bot::State::new());
    let bot_builder = {
        if let Some(format_script) = args.format_script {
            bot_builder
                .with_format_script(format_script)
                .unwrap_or_else(|err| {
                    error!("Failed to set format script: {:?}", err);
                    std::process::exit(1);
                })
        } else {
            bot_builder
        }
    };
    let (stdin_lines_sender, stdin_lines) = channel(1);
    std::thread::spawn(move || {
        stream::iter_ok::<_, ()>(
            BufReader::new(std::io::stdin())
                .lines()
                .filter_map(Result::ok),
        )
        .forward(stdin_lines_sender.sink_map_err(|e| error!("sink error: {}", e)))
        .wait()
    });

    // If we have an email, send an enable message.
    let maybe_enable_message = stream::iter_ok(args.email.as_ref().map(|_| "enable".to_string()));
    let stdin_lines = maybe_enable_message.chain(stdin_lines);

    let email = args.email.clone();
    let use_json = args.json;

    let message_from_line = move |line: String| {
        if use_json {
            serde_json::from_str::<SimpleInputMessage>(&line)
                .map_err(|e| error!("failed to parse message: {}", e))
                .map(Into::into)
                .ok()
        } else if let Some(email) = &email {
            // If we have an email, send each line from this email.
            Some(spark::Message::test_message(
                spark::Email::new(email.clone()),
                spark::PersonId::new(email.clone()),
                line,
            ))
        } else {
            // If no email was given, parse it from each line.
            lazy_static! {
                static ref LINE_REGEX: Regex =
                    Regex::new(r"^(?P<email>.*): ?(?P<message>.*)$").unwrap();
            };

            LINE_REGEX
                .captures(&line)
                .map(|captures| {
                    let email = captures.name("email").unwrap().as_str();
                    let message = captures.name("message").unwrap().as_str();
                    spark::Message::test_message(
                        spark::Email::new(email.to_string()),
                        spark::PersonId::new(email.to_string()),
                        message.to_string(),
                    )
                })
                .or_else(|| {
                    warn!(r#"input not understood: please send as "<email>: <message>""#);
                    None
                })
        }
    };

    let spark_messages = stdin_lines
        .filter(|line| !line.is_empty())
        .filter_map(message_from_line);
    let spark_client = if use_json {
        ConsoleSparkClient::Json
    } else {
        ConsoleSparkClient::Plain
    };

    let bot = bot_builder.build(gerrit_command_runner, spark_client);
    tokio::run(bot.run(gerrit_event_stream, spark_messages));
}
<commit_msg>add hidden option to change working directory to gerritbot-console<commit_after>// need to raise recursion limit because many combinators
#![recursion_limit = "128"]
#![deny(bare_trait_objects)]

use std::io::{BufRead as _, BufReader, Write as _};
use std::path::PathBuf;

use lazy_static::lazy_static;
use regex::Regex;
use serde::{Deserialize, Serialize};
use serde_json;
use structopt::StructOpt;

use futures::{future, stream, sync::mpsc::channel, Future as _, Sink as _, Stream as _};
use log::{error, info, warn};

use gerritbot as bot;
use gerritbot_gerrit as gerrit;
use gerritbot_spark as spark;

#[derive(StructOpt, Debug)]
#[structopt(name = "gerritbot-console", rename_all = "kebab-case")]
/// Run the gerritbot without actually connecting to Spark. Instead the
/// program's stdin can be used to simulate sending messages. Replies will be
/// sent to stdout. Log messages will only appear on stderr. If the --email
/// option (see below) is given each input line will be treated as message from
/// that user. Otherwise each line has to be prefixed with the email address of
/// the intended sender, a colon (:) and an optional space character.
struct Args {
    /// Gerrit bot username
    ///
    /// This is the username the bot uses to connect to the Gerrit server.
    #[structopt(short, long)]
    username: String,
    /// Gerrit hostname
    ///
    /// Address of the Gerrit server.
    hostname: String,
    /// Gerrit SSH port
    #[structopt(short, long, default_value = "29418")]
    port: u32,
    /// Path to SSH private key
    #[structopt(short, long, parse(from_os_str))]
    identity_file: PathBuf,
    /// Enable verbose output
    #[structopt(short, long)]
    verbose: bool,
    /// Log only errors
    #[structopt(short, long, conflicts_with = "verbose")]
    quiet: bool,
    /// User email address
    ///
    /// If given input messages will be treated as if coming from this user.
    /// Additionally, an "enable" message will be injected before the first
    /// input line.
    ///
    /// This option is mutually exclusive with the --json option.
    #[structopt(short, long, conflicts_with = "json")]
    email: Option<String>,
    /// JSON mode
    ///
    /// Switch input and output to JSON. This is option mutually exclusive with
    /// the --email option.
    #[structopt(long)]
    json: bool,
    /// Lua formatting script
    ///
    /// Can be used to change or test the formatting of messages. If not present
    /// the internal default format script will be used.
    #[structopt(long = "format-script")]
    format_script: Option<String>,
    #[structopt(short = "C", hidden = true)]
    working_directory: Option<PathBuf>,
}

#[derive(Deserialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
struct SimpleInputMessage {
    person_email: spark::Email,
    person_id: Option<spark::PersonId>,
    text: String,
}

fn email_to_person_id(email: &spark::Email) -> spark::PersonId {
    spark::PersonId::new(email.to_string())
}

impl Into<spark::Message> for SimpleInputMessage {
    fn into(self) -> spark::Message {
        let SimpleInputMessage {
            person_email,
            person_id,
            text,
        } = self;
        let person_id = person_id.unwrap_or_else(|| email_to_person_id(&person_email));
        spark::Message::test_message(person_email, person_id, text)
    }
}

#[derive(Serialize, Debug, Clone)]
#[serde(rename_all = "camelCase")]
struct SimpleOutputMessage {
    person_id: spark::PersonId,
    text: String,
}

#[derive(Clone)]
enum ConsoleSparkClient {
    Plain,
    Json,
}

impl bot::SparkClient for ConsoleSparkClient {
    type ReplyFuture = future::FutureResult<(), spark::Error>;
    fn reply(&self, person_id: &spark::PersonId, msg: &str) -> Self::ReplyFuture {
        // Write synchronously and crash if writing fails. There's no point in
        // error handling here.
        match self {
            ConsoleSparkClient::Plain => write!(std::io::stdout(), "{}: {}\n", person_id, msg)
                .expect("writing to stdout failed"),
            ConsoleSparkClient::Json => {
                let message = SimpleOutputMessage {
                    person_id: person_id.clone(),
                    text: msg.to_string(),
                };
                serde_json::to_writer(std::io::stdout(), &message)
                    .expect("writing JSON to stdout failed");
                std::io::stdout()
                    .write(b"\n")
                    .expect("writing to stdout failed");
            }
        }
        future::ok(())
    }
}

fn main() {
    let args = Args::from_args();
    stderrlog::new()
        .module(module_path!())
        .module("gerritbot")
        .module("gerritbot_gerrit")
        .timestamp(stderrlog::Timestamp::Second)
        .verbosity(match (args.quiet, args.verbose) {
            (true, _) => 0,      // ERROR
            (false, false) => 2, // INFO
            (_, true) => 4,      // TRACE
        })
        .init()
        .unwrap();

    if let Some(working_directory) = args.working_directory.as_ref() {
        info!("Changing current directory to {:?}", working_directory);
        std::env::set_current_dir(working_directory).expect("failed to set working directory");
    }

    let connect_to_gerrit = || {
        info!(
            "Connecting to gerrit with username {} at {}:{}",
            args.username, args.hostname, args.port,
        );
        gerrit::Connection::connect(
            format!("{}:{}", args.hostname, args.port),
            args.username.clone(),
            args.identity_file.clone(),
        )
        .unwrap_or_else(|e| {
            error!("failed to connect to gerrit: {}", e);
            std::process::exit(1);
        })
    };
    let gerrit_event_stream = gerrit::extended_event_stream(
        connect_to_gerrit(),
        connect_to_gerrit(),
        bot::request_extended_gerrit_info,
    );
    let gerrit_command_runner = gerrit::CommandRunner::new(connect_to_gerrit());
    let bot_builder = bot::Builder::new(bot::State::new());
    let bot_builder = {
        if let Some(format_script) = args.format_script {
            bot_builder
                .with_format_script(format_script)
                .unwrap_or_else(|err| {
                    error!("Failed to set format script: {:?}", err);
                    std::process::exit(1);
                })
        } else {
            bot_builder
        }
    };
    let (stdin_lines_sender, stdin_lines) = channel(1);
    std::thread::spawn(move || {
        stream::iter_ok::<_, ()>(
            BufReader::new(std::io::stdin())
                .lines()
                .filter_map(Result::ok),
        )
        .forward(stdin_lines_sender.sink_map_err(|e| error!("sink error: {}", e)))
        .wait()
    });

    // If we have an email, send an enable message.
    let maybe_enable_message = stream::iter_ok(args.email.as_ref().map(|_| "enable".to_string()));
    let stdin_lines = maybe_enable_message.chain(stdin_lines);

    let email = args.email.clone();
    let use_json = args.json;

    let message_from_line = move |line: String| {
        if use_json {
            serde_json::from_str::<SimpleInputMessage>(&line)
                .map_err(|e| error!("failed to parse message: {}", e))
                .map(Into::into)
                .ok()
        } else if let Some(email) = &email {
            // If we have an email, send each line from this email.
            Some(spark::Message::test_message(
                spark::Email::new(email.clone()),
                spark::PersonId::new(email.clone()),
                line,
            ))
        } else {
            // If no email was given, parse it from each line.
            lazy_static! {
                static ref LINE_REGEX: Regex =
                    Regex::new(r"^(?P<email>.*): ?(?P<message>.*)$").unwrap();
            };

            LINE_REGEX
                .captures(&line)
                .map(|captures| {
                    let email = captures.name("email").unwrap().as_str();
                    let message = captures.name("message").unwrap().as_str();
                    spark::Message::test_message(
                        spark::Email::new(email.to_string()),
                        spark::PersonId::new(email.to_string()),
                        message.to_string(),
                    )
                })
                .or_else(|| {
                    warn!(r#"input not understood: please send as "<email>: <message>""#);
                    None
                })
        }
    };

    let spark_messages = stdin_lines
        .filter(|line| !line.is_empty())
        .filter_map(message_from_line);
    let spark_client = if use_json {
        ConsoleSparkClient::Json
    } else {
        ConsoleSparkClient::Plain
    };

    let bot = bot_builder.build(gerrit_command_runner, spark_client);
    tokio::run(bot.run(gerrit_event_stream, spark_messages));
}
<|endoftext|>


<commit_before>use crate::i18n::LocalText;
use git2::Repository;
use std::{env, error, path, result};

type Result<T> = result::Result<T, Box<dyn error::Error>>;

/// Dump what we know about the repo
pub fn run() -> Result<()> {
    crate::header("status-header");
    is_setup()?;
    Ok(())
}

/// Evaluate whether this project is pis_setup()?roperly configured
pub fn is_setup() -> Result<bool> {
    let mut checks: Vec<bool> = Vec::new();
    let is_repo = is_repo(env::current_dir()?)?;
    checks.push(is_repo);
    let ret = checks.iter().all(|&v| v);
    eprintln!(
        "{}",
        LocalText::new(if ret { "status-good" } else { "status-bad" }).fmt()
    );
    Ok(ret)
}

/// Are we in a git repo?
pub fn is_repo(path: path::PathBuf) -> Result<bool> {
    let ret = Repository::discover(path).is_ok();
    eprintln!(
        "{} {}",
        LocalText::new("status-is-repo").fmt(),
        LocalText::new(if ret { "status-true" } else { "status-false" }).fmt()
    );
    Ok(ret)
}
<commit_msg>chore(cli): Process status checks in parallel threads<commit_after>use crate::i18n::LocalText;
use git2::Repository;
use std::sync::{Arc, RwLock};
use std::{env, error, path, result, thread};

type Result<T> = result::Result<T, Box<dyn error::Error>>;

/// Dump what we know about the repo
pub fn run() -> Result<()> {
    crate::header("status-header");
    is_setup()?;
    Ok(())
}

/// Evaluate whether this project is pis_setup()?roperly configured
pub fn is_setup() -> Result<bool> {
    let results = Arc::new(RwLock::new(Vec::new()));
    rayon::scope(|s| {
        s.spawn(|_| {
            let repo = is_repo(&env::current_dir().unwrap()).unwrap();
            results.write().unwrap().push(repo);
        });
    });
    let ret = results.read().unwrap().iter().all(|&v| v);
    eprintln!(
        "{}",
        LocalText::new(if ret { "status-good" } else { "status-bad" }).fmt()
    );
    Ok(ret)
}

/// Are we in a git repo?
pub fn is_repo(path: &path::PathBuf) -> Result<bool> {
    let ret = Repository::discover(path).is_ok();
    eprintln!(
        "{} {}",
        LocalText::new("status-is-repo").fmt(),
        LocalText::new(if ret { "status-true" } else { "status-false" }).fmt()
    );
    Ok(ret)
}
<|endoftext|>


<commit_before><commit_msg>Implement bigint AddAssign<commit_after><|endoftext|>