<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

/// Gecko's pseudo-element definition.
#[derive(Clone, Debug, Eq, Hash, MallocSizeOf, PartialEq, ToShmem)]
pub enum PseudoElement {
    % for pseudo in PSEUDOS:
        /// ${pseudo.value}
        % if pseudo.is_tree_pseudo_element():
        ${pseudo.capitalized_pseudo()}(ThinBoxedSlice<Atom>),
        % else:
        ${pseudo.capitalized_pseudo()},
        % endif
    % endfor
    /// ::-webkit-* that we don't recognize
    /// https://github.com/whatwg/compat/issues/103
    UnknownWebkit(Atom),
}

/// Important: If you change this, you should also update Gecko's
/// nsCSSPseudoElements::IsEagerlyCascadedInServo.
<% EAGER_PSEUDOS = ["Before", "After", "FirstLine", "FirstLetter"] %>
<% TREE_PSEUDOS = [pseudo for pseudo in PSEUDOS if pseudo.is_tree_pseudo_element()] %>
<% SIMPLE_PSEUDOS = [pseudo for pseudo in PSEUDOS if not pseudo.is_tree_pseudo_element()] %>

/// The number of eager pseudo-elements.
pub const EAGER_PSEUDO_COUNT: usize = ${len(EAGER_PSEUDOS)};

/// The number of non-functional pseudo-elements.
pub const SIMPLE_PSEUDO_COUNT: usize = ${len(SIMPLE_PSEUDOS)};

/// The number of tree pseudo-elements.
pub const TREE_PSEUDO_COUNT: usize = ${len(TREE_PSEUDOS)};

/// The number of all pseudo-elements.
pub const PSEUDO_COUNT: usize = ${len(PSEUDOS)};

/// The list of eager pseudos.
pub const EAGER_PSEUDOS: [PseudoElement; EAGER_PSEUDO_COUNT] = [
    % for eager_pseudo_name in EAGER_PSEUDOS:
    PseudoElement::${eager_pseudo_name},
    % endfor
];

<%def name="pseudo_element_variant(pseudo, tree_arg='..')">\
PseudoElement::${pseudo.capitalized_pseudo()}${"({})".format(tree_arg) if pseudo.is_tree_pseudo_element() else ""}\
</%def>

impl PseudoElement {
    /// Returns an index of the pseudo-element.
    #[inline]
    pub fn index(&self) -> usize {
        match *self {
            % for i, pseudo in enumerate(PSEUDOS):
            ${pseudo_element_variant(pseudo)} => ${i},
            % endfor
            PseudoElement::UnknownWebkit(..) => unreachable!(),
        }
    }

    /// Returns an array of `None` values.
    ///
    /// FIXME(emilio): Integer generics can't come soon enough.
    pub fn pseudo_none_array<T>() -> [Option<T>; PSEUDO_COUNT] {
        [
            ${",\n            ".join(["None" for pseudo in PSEUDOS])}
        ]
    }

    /// Whether this pseudo-element is an anonymous box.
    #[inline]
    pub fn is_anon_box(&self) -> bool {
        match *self {
            % for pseudo in PSEUDOS:
                % if pseudo.is_anon_box():
                    ${pseudo_element_variant(pseudo)} => true,
                % endif
            % endfor
            _ => false,
        }
    }

    /// Whether this pseudo-element is eagerly-cascaded.
    #[inline]
    pub fn is_eager(&self) -> bool {
        matches!(*self,
                 ${" | ".join(map(lambda name: "PseudoElement::{}".format(name), EAGER_PSEUDOS))})
    }

    /// Whether this pseudo-element is tree pseudo-element.
    #[inline]
    pub fn is_tree_pseudo_element(&self) -> bool {
        match *self {
            % for pseudo in TREE_PSEUDOS:
            ${pseudo_element_variant(pseudo)} => true,
            % endfor
            _ => false,
        }
    }

    /// Whether this pseudo-element is an unknown Webkit-prefixed pseudo-element.
    #[inline]
    pub fn is_unknown_webkit_pseudo_element(&self) -> bool {
        matches!(*self, PseudoElement::UnknownWebkit(..))
    }

    /// Gets the flags associated to this pseudo-element, or 0 if it's an
    /// anonymous box.
    pub fn flags(&self) -> u32 {
        match *self {
            % for pseudo in PSEUDOS:
                ${pseudo_element_variant(pseudo)} =>
                % if pseudo.is_tree_pseudo_element():
                    structs::CSS_PSEUDO_ELEMENT_ENABLED_IN_UA_SHEETS_AND_CHROME,
                % elif pseudo.is_anon_box():
                    structs::CSS_PSEUDO_ELEMENT_ENABLED_IN_UA_SHEETS,
                % else:
                    structs::SERVO_CSS_PSEUDO_ELEMENT_FLAGS_${pseudo.pseudo_ident},
                % endif
            % endfor
            PseudoElement::UnknownWebkit(..) => 0,
        }
    }

    /// Construct a pseudo-element from a `PseudoStyleType`.
    #[inline]
    pub fn from_pseudo_type(type_: PseudoStyleType) -> Option<Self> {
        match type_ {
            % for pseudo in PSEUDOS:
            % if not pseudo.is_tree_pseudo_element():
                PseudoStyleType::${pseudo.pseudo_ident} => {
                    Some(${pseudo_element_variant(pseudo)})
                },
            % endif
            % endfor
            _ => None,
        }
    }

    /// Construct a `PseudoStyleType` from a pseudo-element
    #[inline]
    pub fn pseudo_type(&self) -> PseudoStyleType {
        match *self {
            % for pseudo in PSEUDOS:
            % if pseudo.is_tree_pseudo_element():
                PseudoElement::${pseudo.capitalized_pseudo()}(..) => PseudoStyleType::XULTree,
            % else:
                PseudoElement::${pseudo.capitalized_pseudo()} => PseudoStyleType::${pseudo.pseudo_ident},
            % endif
            % endfor
            PseudoElement::UnknownWebkit(..) => unreachable!(),
        }
    }

    /// Get the argument list of a tree pseudo-element.
    #[inline]
    pub fn tree_pseudo_args(&self) -> Option<<&[Atom]> {
        match *self {
            % for pseudo in TREE_PSEUDOS:
            PseudoElement::${pseudo.capitalized_pseudo()}(ref args) => Some(args),
            % endfor
            _ => None,
        }
    }

    /// Construct a tree pseudo-element from atom and args.
    #[inline]
    pub fn from_tree_pseudo_atom(atom: &Atom, args: Box<[Atom]>) -> Option<Self> {
        % for pseudo in PSEUDOS:
        % if pseudo.is_tree_pseudo_element():
            if atom == &atom!("${pseudo.value}") {
                return Some(PseudoElement::${pseudo.capitalized_pseudo()}(args.into()));
            }
        % endif
        % endfor
        None
    }

    /// Constructs a pseudo-element from a string of text.
    ///
    /// Returns `None` if the pseudo-element is not recognised.
    #[inline]
    pub fn from_slice(name: &str) -> Option<Self> {
        // We don't need to support tree pseudos because functional
        // pseudo-elements needs arguments, and thus should be created
        // via other methods.
        match_ignore_ascii_case! { name,
            % for pseudo in SIMPLE_PSEUDOS:
            "${pseudo.value[1:]}" => {
                return Some(${pseudo_element_variant(pseudo)})
            },
            % endfor
            // Alias some legacy prefixed pseudos to their standardized name at parse time:
            "-moz-selection" => {
                return Some(PseudoElement::Selection);
            },
            "-moz-placeholder" => {
                return Some(PseudoElement::Placeholder);
            },
            "-moz-list-bullet" | "-moz-list-number" => {
                return Some(PseudoElement::Marker);
            },
            _ => {
                if starts_with_ignore_ascii_case(name, "-moz-tree-") {
                    return PseudoElement::tree_pseudo_element(name, Box::new([]))
                }
                if static_prefs::pref!("layout.css.unknown-webkit-pseudo-element") {
                    const WEBKIT_PREFIX: &str = "-webkit-";
                    if starts_with_ignore_ascii_case(name, WEBKIT_PREFIX) {
                        let part = string_as_ascii_lowercase(&name[WEBKIT_PREFIX.len()..]);
                        return Some(PseudoElement::UnknownWebkit(part.into()));
                    }
                }
            }
        }

        None
    }

    /// Constructs a tree pseudo-element from the given name and arguments.
    /// "name" must start with "-moz-tree-".
    ///
    /// Returns `None` if the pseudo-element is not recognized.
    #[inline]
    pub fn tree_pseudo_element(name: &str, args: Box<[Atom]>) -> Option<Self> {
        debug_assert!(starts_with_ignore_ascii_case(name, "-moz-tree-"));
        let tree_part = &name[10..];
        % for pseudo in TREE_PSEUDOS:
            if tree_part.eq_ignore_ascii_case("${pseudo.value[11:]}") {
                return Some(${pseudo_element_variant(pseudo, "args.into()")});
            }
        % endfor
        None
    }
}

impl ToCss for PseudoElement {
    fn to_css<W>(&self, dest: &mut W) -> fmt::Result where W: fmt::Write {
        dest.write_char(':')?;
        match *self {
            % for pseudo in PSEUDOS:
                ${pseudo_element_variant(pseudo)} => dest.write_str("${pseudo.value}")?,
            % endfor
            PseudoElement::UnknownWebkit(ref atom) => {
                dest.write_str(":-webkit-")?;
                serialize_atom_identifier(atom, dest)?;
            }
        }
        if let Some(args) = self.tree_pseudo_args() {
            if !args.is_empty() {
                dest.write_char('(')?;
                let mut iter = args.iter();
                if let Some(first) = iter.next() {
                    serialize_atom_identifier(&first, dest)?;
                    for item in iter {
                        dest.write_str(", ")?;
                        serialize_atom_identifier(item, dest)?;
                    }
                }
                dest.write_char(')')?;
            }
        }
        Ok(())
    }
}
<commit_msg>style: Remove layout.css.unknown-webkit-pseudo-element.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

/// Gecko's pseudo-element definition.
#[derive(Clone, Debug, Eq, Hash, MallocSizeOf, PartialEq, ToShmem)]
pub enum PseudoElement {
    % for pseudo in PSEUDOS:
        /// ${pseudo.value}
        % if pseudo.is_tree_pseudo_element():
        ${pseudo.capitalized_pseudo()}(ThinBoxedSlice<Atom>),
        % else:
        ${pseudo.capitalized_pseudo()},
        % endif
    % endfor
    /// ::-webkit-* that we don't recognize
    /// https://github.com/whatwg/compat/issues/103
    UnknownWebkit(Atom),
}

/// Important: If you change this, you should also update Gecko's
/// nsCSSPseudoElements::IsEagerlyCascadedInServo.
<% EAGER_PSEUDOS = ["Before", "After", "FirstLine", "FirstLetter"] %>
<% TREE_PSEUDOS = [pseudo for pseudo in PSEUDOS if pseudo.is_tree_pseudo_element()] %>
<% SIMPLE_PSEUDOS = [pseudo for pseudo in PSEUDOS if not pseudo.is_tree_pseudo_element()] %>

/// The number of eager pseudo-elements.
pub const EAGER_PSEUDO_COUNT: usize = ${len(EAGER_PSEUDOS)};

/// The number of non-functional pseudo-elements.
pub const SIMPLE_PSEUDO_COUNT: usize = ${len(SIMPLE_PSEUDOS)};

/// The number of tree pseudo-elements.
pub const TREE_PSEUDO_COUNT: usize = ${len(TREE_PSEUDOS)};

/// The number of all pseudo-elements.
pub const PSEUDO_COUNT: usize = ${len(PSEUDOS)};

/// The list of eager pseudos.
pub const EAGER_PSEUDOS: [PseudoElement; EAGER_PSEUDO_COUNT] = [
    % for eager_pseudo_name in EAGER_PSEUDOS:
    PseudoElement::${eager_pseudo_name},
    % endfor
];

<%def name="pseudo_element_variant(pseudo, tree_arg='..')">\
PseudoElement::${pseudo.capitalized_pseudo()}${"({})".format(tree_arg) if pseudo.is_tree_pseudo_element() else ""}\
</%def>

impl PseudoElement {
    /// Returns an index of the pseudo-element.
    #[inline]
    pub fn index(&self) -> usize {
        match *self {
            % for i, pseudo in enumerate(PSEUDOS):
            ${pseudo_element_variant(pseudo)} => ${i},
            % endfor
            PseudoElement::UnknownWebkit(..) => unreachable!(),
        }
    }

    /// Returns an array of `None` values.
    ///
    /// FIXME(emilio): Integer generics can't come soon enough.
    pub fn pseudo_none_array<T>() -> [Option<T>; PSEUDO_COUNT] {
        [
            ${",\n            ".join(["None" for pseudo in PSEUDOS])}
        ]
    }

    /// Whether this pseudo-element is an anonymous box.
    #[inline]
    pub fn is_anon_box(&self) -> bool {
        match *self {
            % for pseudo in PSEUDOS:
                % if pseudo.is_anon_box():
                    ${pseudo_element_variant(pseudo)} => true,
                % endif
            % endfor
            _ => false,
        }
    }

    /// Whether this pseudo-element is eagerly-cascaded.
    #[inline]
    pub fn is_eager(&self) -> bool {
        matches!(*self,
                 ${" | ".join(map(lambda name: "PseudoElement::{}".format(name), EAGER_PSEUDOS))})
    }

    /// Whether this pseudo-element is tree pseudo-element.
    #[inline]
    pub fn is_tree_pseudo_element(&self) -> bool {
        match *self {
            % for pseudo in TREE_PSEUDOS:
            ${pseudo_element_variant(pseudo)} => true,
            % endfor
            _ => false,
        }
    }

    /// Whether this pseudo-element is an unknown Webkit-prefixed pseudo-element.
    #[inline]
    pub fn is_unknown_webkit_pseudo_element(&self) -> bool {
        matches!(*self, PseudoElement::UnknownWebkit(..))
    }

    /// Gets the flags associated to this pseudo-element, or 0 if it's an
    /// anonymous box.
    pub fn flags(&self) -> u32 {
        match *self {
            % for pseudo in PSEUDOS:
                ${pseudo_element_variant(pseudo)} =>
                % if pseudo.is_tree_pseudo_element():
                    structs::CSS_PSEUDO_ELEMENT_ENABLED_IN_UA_SHEETS_AND_CHROME,
                % elif pseudo.is_anon_box():
                    structs::CSS_PSEUDO_ELEMENT_ENABLED_IN_UA_SHEETS,
                % else:
                    structs::SERVO_CSS_PSEUDO_ELEMENT_FLAGS_${pseudo.pseudo_ident},
                % endif
            % endfor
            PseudoElement::UnknownWebkit(..) => 0,
        }
    }

    /// Construct a pseudo-element from a `PseudoStyleType`.
    #[inline]
    pub fn from_pseudo_type(type_: PseudoStyleType) -> Option<Self> {
        match type_ {
            % for pseudo in PSEUDOS:
            % if not pseudo.is_tree_pseudo_element():
                PseudoStyleType::${pseudo.pseudo_ident} => {
                    Some(${pseudo_element_variant(pseudo)})
                },
            % endif
            % endfor
            _ => None,
        }
    }

    /// Construct a `PseudoStyleType` from a pseudo-element
    #[inline]
    pub fn pseudo_type(&self) -> PseudoStyleType {
        match *self {
            % for pseudo in PSEUDOS:
            % if pseudo.is_tree_pseudo_element():
                PseudoElement::${pseudo.capitalized_pseudo()}(..) => PseudoStyleType::XULTree,
            % else:
                PseudoElement::${pseudo.capitalized_pseudo()} => PseudoStyleType::${pseudo.pseudo_ident},
            % endif
            % endfor
            PseudoElement::UnknownWebkit(..) => unreachable!(),
        }
    }

    /// Get the argument list of a tree pseudo-element.
    #[inline]
    pub fn tree_pseudo_args(&self) -> Option<<&[Atom]> {
        match *self {
            % for pseudo in TREE_PSEUDOS:
            PseudoElement::${pseudo.capitalized_pseudo()}(ref args) => Some(args),
            % endfor
            _ => None,
        }
    }

    /// Construct a tree pseudo-element from atom and args.
    #[inline]
    pub fn from_tree_pseudo_atom(atom: &Atom, args: Box<[Atom]>) -> Option<Self> {
        % for pseudo in PSEUDOS:
        % if pseudo.is_tree_pseudo_element():
            if atom == &atom!("${pseudo.value}") {
                return Some(PseudoElement::${pseudo.capitalized_pseudo()}(args.into()));
            }
        % endif
        % endfor
        None
    }

    /// Constructs a pseudo-element from a string of text.
    ///
    /// Returns `None` if the pseudo-element is not recognised.
    #[inline]
    pub fn from_slice(name: &str) -> Option<Self> {
        // We don't need to support tree pseudos because functional
        // pseudo-elements needs arguments, and thus should be created
        // via other methods.
        match_ignore_ascii_case! { name,
            % for pseudo in SIMPLE_PSEUDOS:
            "${pseudo.value[1:]}" => {
                return Some(${pseudo_element_variant(pseudo)})
            },
            % endfor
            // Alias some legacy prefixed pseudos to their standardized name at parse time:
            "-moz-selection" => {
                return Some(PseudoElement::Selection);
            },
            "-moz-placeholder" => {
                return Some(PseudoElement::Placeholder);
            },
            "-moz-list-bullet" | "-moz-list-number" => {
                return Some(PseudoElement::Marker);
            },
            _ => {
                if starts_with_ignore_ascii_case(name, "-moz-tree-") {
                    return PseudoElement::tree_pseudo_element(name, Box::new([]))
                }
                const WEBKIT_PREFIX: &str = "-webkit-";
                if starts_with_ignore_ascii_case(name, WEBKIT_PREFIX) {
                    let part = string_as_ascii_lowercase(&name[WEBKIT_PREFIX.len()..]);
                    return Some(PseudoElement::UnknownWebkit(part.into()));
                }
            }
        }

        None
    }

    /// Constructs a tree pseudo-element from the given name and arguments.
    /// "name" must start with "-moz-tree-".
    ///
    /// Returns `None` if the pseudo-element is not recognized.
    #[inline]
    pub fn tree_pseudo_element(name: &str, args: Box<[Atom]>) -> Option<Self> {
        debug_assert!(starts_with_ignore_ascii_case(name, "-moz-tree-"));
        let tree_part = &name[10..];
        % for pseudo in TREE_PSEUDOS:
            if tree_part.eq_ignore_ascii_case("${pseudo.value[11:]}") {
                return Some(${pseudo_element_variant(pseudo, "args.into()")});
            }
        % endfor
        None
    }
}

impl ToCss for PseudoElement {
    fn to_css<W>(&self, dest: &mut W) -> fmt::Result where W: fmt::Write {
        dest.write_char(':')?;
        match *self {
            % for pseudo in PSEUDOS:
                ${pseudo_element_variant(pseudo)} => dest.write_str("${pseudo.value}")?,
            % endfor
            PseudoElement::UnknownWebkit(ref atom) => {
                dest.write_str(":-webkit-")?;
                serialize_atom_identifier(atom, dest)?;
            }
        }
        if let Some(args) = self.tree_pseudo_args() {
            if !args.is_empty() {
                dest.write_char('(')?;
                let mut iter = args.iter();
                if let Some(first) = iter.next() {
                    serialize_atom_identifier(&first, dest)?;
                    for item in iter {
                        dest.write_str(", ")?;
                        serialize_atom_identifier(item, dest)?;
                    }
                }
                dest.write_char(')')?;
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;

#[cfg(all(feature = "rustc_ser_type", not(feature = "serde_type")))]
use serialize::json::ToJson;
#[cfg(feature = "serde_type")]
use serde_json::value::ToJson;

use regex::{Regex, Captures};

use template::Template;
use render::{Renderable, RenderError, RenderContext};
use context::Context;
use helpers::{self, HelperDef};
use directives::{self, DirectiveDef};
use support::str::StringWriter;
use error::{TemplateError, TemplateFileError, TemplateRenderError};

lazy_static!{
    static ref DEFAULT_REPLACE: Regex = Regex::new(">|<|\"|&").unwrap();
}

/// This type represents an *escape fn*, that is a function who's purpose it is
/// to escape potentially problematic characters in a string.
///
/// An *escape fn* is represented as a `Box` to avoid unnecessary type
/// parameters (and because traits cannot be aliased using `type`).
pub type EscapeFn = Box<Fn(&str) -> String + Send + Sync>;

/// The default *escape fn* replaces the characters `&"<>`
/// with the equivalent html / xml entities.
pub fn html_escape(data: &str) -> String {
    DEFAULT_REPLACE.replace_all(data, |cap: &Captures| {
        match cap.at(0) {
            Some("<") => "<".to_owned(),
            Some(">") => ">".to_owned(),
            Some("\"") => """.to_owned(),
            Some("&") => "&".to_owned(),
            _ => unreachable!(),
        }
    })
}

/// `EscapeFn` that donot change any thing. Useful when using in a non-html
/// environment.
pub fn no_escape(data: &str) -> String {
    data.to_owned()
}

/// The single entry point of your Handlebars templates
///
/// It maintains compiled templates and registered helpers.
pub struct Registry {
    templates: HashMap<String, Template>,
    helpers: HashMap<String, Box<HelperDef + 'static>>,
    directives: HashMap<String, Box<DirectiveDef + 'static>>,
    escape_fn: EscapeFn,
    source_map: bool,
}

impl Registry {
    pub fn new() -> Registry {
        let r = Registry {
            templates: HashMap::new(),
            helpers: HashMap::new(),
            directives: HashMap::new(),
            escape_fn: Box::new(html_escape),
            source_map: true,
        };

        r.setup_builtins()
    }

    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    fn setup_builtins(mut self) -> Registry {
        self.register_helper("if", Box::new(helpers::IF_HELPER));
        self.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        self.register_helper("each", Box::new(helpers::EACH_HELPER));
        self.register_helper("with", Box::new(helpers::WITH_HELPER));
        self.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        self.register_helper("raw", Box::new(helpers::RAW_HELPER));
        self.register_helper(">", Box::new(helpers::INCLUDE_HELPER));
        self.register_helper("block", Box::new(helpers::BLOCK_HELPER));
        self.register_helper("partial", Box::new(helpers::PARTIAL_HELPER));
        self.register_helper("log", Box::new(helpers::LOG_HELPER));

        self.register_decorator("inline", Box::new(directives::INLINE_DIRECTIVE));
        self
    }

    #[cfg(feature = "partial4")]
    fn setup_builtins(mut self) -> Registry {
        self.register_helper("if", Box::new(helpers::IF_HELPER));
        self.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        self.register_helper("each", Box::new(helpers::EACH_HELPER));
        self.register_helper("with", Box::new(helpers::WITH_HELPER));
        self.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        self.register_helper("raw", Box::new(helpers::RAW_HELPER));
        self.register_helper(">", Box::new(helpers::INCLUDE_HELPER));
        self.register_helper("log", Box::new(helpers::LOG_HELPER));

        self.register_decorator("inline", Box::new(directives::INLINE_DIRECTIVE));
        self
    }

    /// Enable handlebars template source map
    pub fn source_map_enable(&mut self, enable: bool) {
        self.source_map = enable;
    }

    /// Register a template
    pub fn register_template(&mut self, name: &str, mut template: Template) {
        template.name = Some(name.to_owned());
        self.templates.insert(name.to_string(), template);
    }

    /// Register a template string
    pub fn register_template_string(&mut self,
                                    name: &str,
                                    tpl_str: String)
                                    -> Result<(), TemplateError> {
        try!(Template::compile_with_name(tpl_str, name.to_owned(), self.source_map)
                 .and_then(|t| Ok(self.templates.insert(name.to_string(), t))));
        Ok(())
    }

    /// Register a template from a path
    pub fn register_template_file<P: AsRef<Path>>(&mut self,
                                                  name: &str,
                                                  tpl_path: P)
                                                  -> Result<(), TemplateFileError> {
        let mut file = try!(File::open(tpl_path)
                                .map_err(|e| TemplateFileError::IOError(e, name.to_owned())));
        self.register_template_source(name, &mut file)
    }

    /// Register a template from `std::io::Read` source
    pub fn register_template_source(&mut self,
                                    name: &str,
                                    tpl_source: &mut Read)
                                    -> Result<(), TemplateFileError> {
        let mut buf = String::new();
        try!(tpl_source.read_to_string(&mut buf)
                       .map_err(|e| TemplateFileError::IOError(e, name.to_owned())));
        try!(self.register_template_string(name, buf));
        Ok(())
    }

    /// remove a template from the registry
    pub fn unregister_template(&mut self, name: &str) {
        self.templates.remove(name);
    }

    /// register a helper
    pub fn register_helper(&mut self,
                           name: &str,
                           def: Box<HelperDef + 'static>)
                           -> Option<Box<HelperDef + 'static>> {
        self.helpers.insert(name.to_string(), def)
    }

    /// register a directive
    pub fn register_decorator(&mut self,
                              name: &str,
                              def: Box<DirectiveDef + 'static>)
                              -> Option<Box<DirectiveDef + 'static>> {
        self.directives.insert(name.to_string(), def)
    }

    /// Register a new *escape fn* to be used from now on by this registry.
    pub fn register_escape_fn<F: 'static + Fn(&str) -> String + Send + Sync>(&mut self,
                                                                             escape_fn: F) {
        self.escape_fn = Box::new(escape_fn);
    }

    /// Restore the default *escape fn*.
    pub fn unregister_escape_fn(&mut self) {
        self.escape_fn = Box::new(html_escape);
    }

    /// Get a reference to the current *escape fn*.
    pub fn get_escape_fn(&self) -> &Fn(&str) -> String {
        &*self.escape_fn
    }

    /// Return a registered template,
    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    /// Return a registered helper
    pub fn get_helper(&self, name: &str) -> Option<&Box<HelperDef + 'static>> {
        self.helpers.get(name)
    }

    /// Return a registered directive, aka decorator
    pub fn get_decorator(&self, name: &str) -> Option<&Box<DirectiveDef + 'static>> {
        self.directives.get(name)
    }

    /// Return all templates registered
    pub fn get_templates(&self) -> &HashMap<String, Template> {
        &self.templates
    }

    /// Unregister all templates
    pub fn clear_templates(&mut self) {
        self.templates.clear();
    }

    /// Render a registered template with some data into a string
    pub fn render<T>(&self, name: &str, ctx: &T) -> Result<String, RenderError>
        where T: ToJson
    {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.renderw(name, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// Render a registered template with some data to the `std::io::Write`
    pub fn renderw(&self,
                   name: &str,
                   context: &Context,
                   writer: &mut Write)
                   -> Result<(), RenderError> {
        let template = self.get_template(&name.to_string());

        if let Some(t) = template {
            let mut render_context = RenderContext::new(writer);
            render_context.root_template = t.name.clone();
            (*t).render(context, self, &mut render_context)
        } else {
            Err(RenderError::new(format!("Template not found: {}", name)))
        }
    }

    /// render a template string using current registry without register it
    pub fn template_render<T>(&self,
                              template_string: &str,
                              ctx: &T)
                              -> Result<String, TemplateRenderError>
        where T: ToJson
    {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.template_renderw(template_string, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// render a template string using current registry without register it
    pub fn template_renderw(&self,
                            template_string: &str,
                            context: &Context,
                            writer: &mut Write)
                            -> Result<(), TemplateRenderError> {
        let tpl = try!(Template::compile(template_string));
        let mut render_context = RenderContext::new(writer);
        tpl.render(context, self, &mut render_context).map_err(TemplateRenderError::from)
    }

    /// render a template source using current registry without register it
    pub fn template_renderw2(&self,
                             template_source: &mut Read,
                             context: &Context,
                             writer: &mut Write)
                             -> Result<(), TemplateRenderError> {
        let mut tpl_str = String::new();
        try!(template_source.read_to_string(&mut tpl_str).map_err(|e| {
            TemplateRenderError::IOError(e, "Unamed template source".to_owned())
        }));
        self.template_renderw(&tpl_str, context, writer)
    }
}

#[cfg(test)]
mod test {
    use template::Template;
    use registry::Registry;
    use render::{RenderContext, Renderable, RenderError, Helper};
    use helpers::HelperDef;
    use context::Context;
    use support::str::StringWriter;
    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    use error::TemplateRenderError;

    #[derive(Clone, Copy)]
    struct DummyHelper;

    impl HelperDef for DummyHelper {
        fn call(&self,
                c: &Context,
                h: &Helper,
                r: &Registry,
                rc: &mut RenderContext)
                -> Result<(), RenderError> {
            try!(h.template().unwrap().render(c, r, rc));
            Ok(())
        }
    }

    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    static DUMMY_HELPER: DummyHelper = DummyHelper;

    #[test]
    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    fn test_registry_operations() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let t2 = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index2", t2.clone());

        assert_eq!(r.get_template("index").unwrap().elements, t.elements);
        assert_eq!(r.templates.len(), 2);

        r.unregister_template("index");
        assert_eq!(r.templates.len(), 1);

        r.clear_templates();
        assert_eq!(r.templates.len(), 0);

        r.register_helper("dummy", Box::new(DUMMY_HELPER));

        // built-in helpers plus 1
        assert_eq!(r.helpers.len(), 10 + 1);
    }

    #[test]
    fn test_renderw() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.renderw("index", &context, &mut sw).ok().unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

    }

    #[test]
    fn test_escape_fn() {
        let mut r = Registry::new();

        let input = String::from("\"<>&");

        r.register_template_string("test", String::from("{{this}}")).unwrap();

        assert_eq!(""<>&", r.render("test", &input).unwrap());

        r.register_escape_fn(|s| s.into());

        assert_eq!("\"<>&", r.render("test", &input).unwrap());

        r.unregister_escape_fn();

        assert_eq!(""<>&", r.render("test", &input).unwrap());
    }

    #[test]
    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    fn test_template_render() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        assert_eq!("

".to_string(),
                   r.template_render("{{> index}}", &{}).unwrap());

        assert_eq!("hello world".to_string(),
                   r.template_render("hello {{this}}", &"world".to_string()).unwrap());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.template_renderw("{{> index}}", &context, &mut sw).unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

        // fail for template error
        match r.template_render("{{ hello", &{}).unwrap_err() {
            TemplateRenderError::TemplateError(_) => {}
            _ => {
                panic!();
            }
        }

        // fail to render error
        match r.template_render("{{> notfound}}", &{}).unwrap_err() {
            TemplateRenderError::RenderError(_) => {}
            _ => {
                panic!();
            }
        }
    }
}
<commit_msg>(feat) added a new public API: register_partial<commit_after>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;

#[cfg(all(feature = "rustc_ser_type", not(feature = "serde_type")))]
use serialize::json::ToJson;
#[cfg(feature = "serde_type")]
use serde_json::value::ToJson;

use regex::{Regex, Captures};

use template::Template;
use render::{Renderable, RenderError, RenderContext};
use context::Context;
use helpers::{self, HelperDef};
use directives::{self, DirectiveDef};
use support::str::StringWriter;
use error::{TemplateError, TemplateFileError, TemplateRenderError};

lazy_static!{
    static ref DEFAULT_REPLACE: Regex = Regex::new(">|<|\"|&").unwrap();
}

/// This type represents an *escape fn*, that is a function who's purpose it is
/// to escape potentially problematic characters in a string.
///
/// An *escape fn* is represented as a `Box` to avoid unnecessary type
/// parameters (and because traits cannot be aliased using `type`).
pub type EscapeFn = Box<Fn(&str) -> String + Send + Sync>;

/// The default *escape fn* replaces the characters `&"<>`
/// with the equivalent html / xml entities.
pub fn html_escape(data: &str) -> String {
    DEFAULT_REPLACE.replace_all(data, |cap: &Captures| {
        match cap.at(0) {
            Some("<") => "<".to_owned(),
            Some(">") => ">".to_owned(),
            Some("\"") => """.to_owned(),
            Some("&") => "&".to_owned(),
            _ => unreachable!(),
        }
    })
}

/// `EscapeFn` that donot change any thing. Useful when using in a non-html
/// environment.
pub fn no_escape(data: &str) -> String {
    data.to_owned()
}

/// The single entry point of your Handlebars templates
///
/// It maintains compiled templates and registered helpers.
pub struct Registry {
    templates: HashMap<String, Template>,
    helpers: HashMap<String, Box<HelperDef + 'static>>,
    directives: HashMap<String, Box<DirectiveDef + 'static>>,
    escape_fn: EscapeFn,
    source_map: bool,
}

impl Registry {
    pub fn new() -> Registry {
        let r = Registry {
            templates: HashMap::new(),
            helpers: HashMap::new(),
            directives: HashMap::new(),
            escape_fn: Box::new(html_escape),
            source_map: true,
        };

        r.setup_builtins()
    }

    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    fn setup_builtins(mut self) -> Registry {
        self.register_helper("if", Box::new(helpers::IF_HELPER));
        self.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        self.register_helper("each", Box::new(helpers::EACH_HELPER));
        self.register_helper("with", Box::new(helpers::WITH_HELPER));
        self.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        self.register_helper("raw", Box::new(helpers::RAW_HELPER));
        self.register_helper(">", Box::new(helpers::INCLUDE_HELPER));
        self.register_helper("block", Box::new(helpers::BLOCK_HELPER));
        self.register_helper("partial", Box::new(helpers::PARTIAL_HELPER));
        self.register_helper("log", Box::new(helpers::LOG_HELPER));

        self.register_decorator("inline", Box::new(directives::INLINE_DIRECTIVE));
        self
    }

    #[cfg(feature = "partial4")]
    fn setup_builtins(mut self) -> Registry {
        self.register_helper("if", Box::new(helpers::IF_HELPER));
        self.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        self.register_helper("each", Box::new(helpers::EACH_HELPER));
        self.register_helper("with", Box::new(helpers::WITH_HELPER));
        self.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        self.register_helper("raw", Box::new(helpers::RAW_HELPER));
        self.register_helper(">", Box::new(helpers::INCLUDE_HELPER));
        self.register_helper("log", Box::new(helpers::LOG_HELPER));

        self.register_decorator("inline", Box::new(directives::INLINE_DIRECTIVE));
        self
    }

    /// Enable handlebars template source map
    pub fn source_map_enable(&mut self, enable: bool) {
        self.source_map = enable;
    }

    /// Register a template
    pub fn register_template(&mut self, name: &str, mut template: Template) {
        template.name = Some(name.to_owned());
        self.templates.insert(name.to_string(), template);
    }

    /// Register a template string
    pub fn register_template_string(&mut self,
                                    name: &str,
                                    tpl_str: String)
                                    -> Result<(), TemplateError> {
        try!(Template::compile_with_name(tpl_str, name.to_owned(), self.source_map)
                 .and_then(|t| Ok(self.templates.insert(name.to_string(), t))));
        Ok(())
    }

    /// Register a partial string
    pub fn register_partial(&mut self,
                            name: &str,
                            partial_str: String)
                            -> Result<(), TemplateError> {
        self.register_template_string(name, partial_str)
    }

    /// Register a template from a path
    pub fn register_template_file<P: AsRef<Path>>(&mut self,
                                                  name: &str,
                                                  tpl_path: P)
                                                  -> Result<(), TemplateFileError> {
        let mut file = try!(File::open(tpl_path)
                                .map_err(|e| TemplateFileError::IOError(e, name.to_owned())));
        self.register_template_source(name, &mut file)
    }

    /// Register a template from `std::io::Read` source
    pub fn register_template_source(&mut self,
                                    name: &str,
                                    tpl_source: &mut Read)
                                    -> Result<(), TemplateFileError> {
        let mut buf = String::new();
        try!(tpl_source.read_to_string(&mut buf)
                       .map_err(|e| TemplateFileError::IOError(e, name.to_owned())));
        try!(self.register_template_string(name, buf));
        Ok(())
    }

    /// remove a template from the registry
    pub fn unregister_template(&mut self, name: &str) {
        self.templates.remove(name);
    }

    /// register a helper
    pub fn register_helper(&mut self,
                           name: &str,
                           def: Box<HelperDef + 'static>)
                           -> Option<Box<HelperDef + 'static>> {
        self.helpers.insert(name.to_string(), def)
    }

    /// register a directive
    pub fn register_decorator(&mut self,
                              name: &str,
                              def: Box<DirectiveDef + 'static>)
                              -> Option<Box<DirectiveDef + 'static>> {
        self.directives.insert(name.to_string(), def)
    }

    /// Register a new *escape fn* to be used from now on by this registry.
    pub fn register_escape_fn<F: 'static + Fn(&str) -> String + Send + Sync>(&mut self,
                                                                             escape_fn: F) {
        self.escape_fn = Box::new(escape_fn);
    }

    /// Restore the default *escape fn*.
    pub fn unregister_escape_fn(&mut self) {
        self.escape_fn = Box::new(html_escape);
    }

    /// Get a reference to the current *escape fn*.
    pub fn get_escape_fn(&self) -> &Fn(&str) -> String {
        &*self.escape_fn
    }

    /// Return a registered template,
    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    /// Return a registered helper
    pub fn get_helper(&self, name: &str) -> Option<&Box<HelperDef + 'static>> {
        self.helpers.get(name)
    }

    /// Return a registered directive, aka decorator
    pub fn get_decorator(&self, name: &str) -> Option<&Box<DirectiveDef + 'static>> {
        self.directives.get(name)
    }

    /// Return all templates registered
    pub fn get_templates(&self) -> &HashMap<String, Template> {
        &self.templates
    }

    /// Unregister all templates
    pub fn clear_templates(&mut self) {
        self.templates.clear();
    }

    /// Render a registered template with some data into a string
    pub fn render<T>(&self, name: &str, ctx: &T) -> Result<String, RenderError>
        where T: ToJson
    {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.renderw(name, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// Render a registered template with some data to the `std::io::Write`
    pub fn renderw(&self,
                   name: &str,
                   context: &Context,
                   writer: &mut Write)
                   -> Result<(), RenderError> {
        let template = self.get_template(&name.to_string());

        if let Some(t) = template {
            let mut render_context = RenderContext::new(writer);
            render_context.root_template = t.name.clone();
            (*t).render(context, self, &mut render_context)
        } else {
            Err(RenderError::new(format!("Template not found: {}", name)))
        }
    }

    /// render a template string using current registry without register it
    pub fn template_render<T>(&self,
                              template_string: &str,
                              ctx: &T)
                              -> Result<String, TemplateRenderError>
        where T: ToJson
    {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.template_renderw(template_string, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// render a template string using current registry without register it
    pub fn template_renderw(&self,
                            template_string: &str,
                            context: &Context,
                            writer: &mut Write)
                            -> Result<(), TemplateRenderError> {
        let tpl = try!(Template::compile(template_string));
        let mut render_context = RenderContext::new(writer);
        tpl.render(context, self, &mut render_context).map_err(TemplateRenderError::from)
    }

    /// render a template source using current registry without register it
    pub fn template_renderw2(&self,
                             template_source: &mut Read,
                             context: &Context,
                             writer: &mut Write)
                             -> Result<(), TemplateRenderError> {
        let mut tpl_str = String::new();
        try!(template_source.read_to_string(&mut tpl_str).map_err(|e| {
            TemplateRenderError::IOError(e, "Unamed template source".to_owned())
        }));
        self.template_renderw(&tpl_str, context, writer)
    }
}

#[cfg(test)]
mod test {
    use template::Template;
    use registry::Registry;
    use render::{RenderContext, Renderable, RenderError, Helper};
    use helpers::HelperDef;
    use context::Context;
    use support::str::StringWriter;
    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    use error::TemplateRenderError;

    #[derive(Clone, Copy)]
    struct DummyHelper;

    impl HelperDef for DummyHelper {
        fn call(&self,
                c: &Context,
                h: &Helper,
                r: &Registry,
                rc: &mut RenderContext)
                -> Result<(), RenderError> {
            try!(h.template().unwrap().render(c, r, rc));
            Ok(())
        }
    }

    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    static DUMMY_HELPER: DummyHelper = DummyHelper;

    #[test]
    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    fn test_registry_operations() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let t2 = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index2", t2.clone());

        assert_eq!(r.get_template("index").unwrap().elements, t.elements);
        assert_eq!(r.templates.len(), 2);

        r.unregister_template("index");
        assert_eq!(r.templates.len(), 1);

        r.clear_templates();
        assert_eq!(r.templates.len(), 0);

        r.register_helper("dummy", Box::new(DUMMY_HELPER));

        // built-in helpers plus 1
        assert_eq!(r.helpers.len(), 10 + 1);
    }

    #[test]
    fn test_renderw() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.renderw("index", &context, &mut sw).ok().unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

    }

    #[test]
    fn test_escape_fn() {
        let mut r = Registry::new();

        let input = String::from("\"<>&");

        r.register_template_string("test", String::from("{{this}}")).unwrap();

        assert_eq!(""<>&", r.render("test", &input).unwrap());

        r.register_escape_fn(|s| s.into());

        assert_eq!("\"<>&", r.render("test", &input).unwrap());

        r.unregister_escape_fn();

        assert_eq!(""<>&", r.render("test", &input).unwrap());
    }

    #[test]
    #[cfg(all(feature="partial_legacy", not(feature="partial4")))]
    fn test_template_render() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        assert_eq!("

".to_string(),
                   r.template_render("{{> index}}", &{}).unwrap());

        assert_eq!("hello world".to_string(),
                   r.template_render("hello {{this}}", &"world".to_string()).unwrap());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.template_renderw("{{> index}}", &context, &mut sw).unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

        // fail for template error
        match r.template_render("{{ hello", &{}).unwrap_err() {
            TemplateRenderError::TemplateError(_) => {}
            _ => {
                panic!();
            }
        }

        // fail to render error
        match r.template_render("{{> notfound}}", &{}).unwrap_err() {
            TemplateRenderError::RenderError(_) => {}
            _ => {
                panic!();
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;

#[cfg(not(feature = "serde_type"))]
use serialize::json::ToJson;
#[cfg(feature = "serde_type")]
use serde::ser::Serialize as ToJson;

use template::{Template};
use render::{Renderable, RenderError, RenderContext};
use helpers::{HelperDef};
use context::{Context};
use helpers;
use support::str::StringWriter;
use error::{TemplateError, TemplateFileError, TemplateRenderError};

/// This type represents an *escape fn*, that is a function who's purpose it is
/// to escape potentially problematic characters in a string.
///
/// The default *escape fn* replaces the characters `&"<>`
/// with the equivalent html / xml entities.
///
/// An *escape fn* is represented as a `Box` to avoid unnecessary type
/// parameters (and because traits cannot be aliased using `type`).
pub type EscapeFn = Box<Fn(&str) -> String + Send + Sync>;

fn get_default_escape_fn() -> EscapeFn {
    Box::new(|data| {
        data.replace("&", "&")
            .replace("\"", """)
            .replace("<", "<")
            .replace(">", ">")
    })
}

/// The single entry point of your Handlebars templates
///
/// It maintains compiled templates and registered helpers.
pub struct Registry {
    templates: HashMap<String, Template>,
    helpers: HashMap<String, Box<HelperDef + 'static>>,
    escape_fn: EscapeFn
}

impl Registry {
    pub fn new() -> Registry {
        let mut r = Registry {
            templates: HashMap::new(),
            helpers: HashMap::new(),
            escape_fn: get_default_escape_fn(),
        };

        r.register_helper("if", Box::new(helpers::IF_HELPER));
        r.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        r.register_helper("each", Box::new(helpers::EACH_HELPER));
        r.register_helper("with", Box::new(helpers::WITH_HELPER));
        r.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        r.register_helper("raw", Box::new(helpers::RAW_HELPER));
        r.register_helper(">", Box::new(helpers::INCLUDE_HELPER));
        r.register_helper("block", Box::new(helpers::BLOCK_HELPER));
        r.register_helper("partial", Box::new(helpers::PARTIAL_HELPER));
        r.register_helper("log", Box::new(helpers::LOG_HELPER));

        r
    }

    /// Register a template
    pub fn register_template(&mut self, name: &str, mut template: Template) {
        template.name = Some(name.to_owned());
        self.templates.insert(name.to_string(), template);
    }

    /// Register a template string
    pub fn register_template_string(&mut self, name: &str, tpl_str: String) -> Result<(), TemplateError> {
        try!(Template::compile_with_name(tpl_str, name.to_owned())
             .and_then(|t| Ok(self.templates.insert(name.to_string(), t))));
        Ok(())
    }

    /// Register a template from a path
    pub fn register_template_file(&mut self, name: &str, tpl_path: &Path) -> Result<(), TemplateFileError> {
        let mut file = try!(File::open(tpl_path));
        self.register_template_source(name, &mut file)
    }

    /// Register a template from `std::io::Read` source
    pub fn register_template_source(&mut self, name: &str, tpl_source: &mut Read) -> Result<(), TemplateFileError> {
        let mut buf = String::new();
        try!(tpl_source.read_to_string(&mut buf));
        try!(Template::compile_with_name(&buf, name.to_owned())
             .and_then(|t| Ok(self.templates.insert(name.to_string(), t))));
        Ok(())
    }

    /// remove a template from the registry
    pub fn unregister_template(&mut self, name: &str) {
        self.templates.remove(name);
    }

    /// register a helper
    pub fn register_helper(&mut self, name: &str, def: Box<HelperDef + 'static>) -> Option<Box<HelperDef + 'static>> {
        self.helpers.insert(name.to_string(), def)
    }

    /// Register a new *escape fn* to be used from now on by this registry.
    pub fn register_escape_fn<F: 'static + Fn(&str) -> String + Send + Sync>(&mut self, escape_fn: F) {
        self.escape_fn = Box::new(escape_fn);
    }

    /// Restore the default *escape fn*.
    pub fn unregister_escape_fn(&mut self) {
        self.escape_fn = get_default_escape_fn();
    }

    /// Get a reference to the current *escape fn*.
    pub fn get_escape_fn(&self) -> &Fn(&str) -> String {
        &*self.escape_fn
    }

    /// Return a registered template,
    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    /// Return a registered helper
    pub fn get_helper(&self, name: &str) -> Option<&Box<HelperDef + 'static>> {
        self.helpers.get(name)
    }

    /// Return all templates registered
    pub fn get_templates(&self) -> &HashMap<String, Template> {
        &self.templates
    }

    /// Unregister all templates
    pub fn clear_templates(&mut self) {
        self.templates.clear();
    }

    /// Render a registered template with some data into a string
    pub fn render<T>(&self, name: &str, ctx: &T) -> Result<String, RenderError> where T: ToJson {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.renderw(name, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// Render a registered template with some data to the `std::io::Write`
    pub fn renderw(&self, name: &str, context: &Context, writer: &mut Write) -> Result<(), RenderError> {
        let template = self.get_template(&name.to_string());

        if let Some(t) = template {
            let mut render_context = RenderContext::new(writer);
            render_context.root_template = t.name.clone();
            (*t).render(context, self, &mut render_context)
        } else {
            Err(RenderError::new(format!("Template not found: {}", name)))
        }
    }

    /// render a template string using current registry without register it
    pub fn template_render<T>(&self, template_string: &str, ctx: &T) -> Result<String, TemplateRenderError> where T: ToJson {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.template_renderw(template_string, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// render a template string using current registry without register it
    pub fn template_renderw(&self, template_string: &str, context: &Context, writer: &mut Write) -> Result<(), TemplateRenderError> {
        let tpl = try!(Template::compile(template_string).map_err(TemplateRenderError::from));
        let mut render_context = RenderContext::new(writer);
        tpl.render(context, self, &mut render_context).map_err(TemplateRenderError::from)
    }
}

#[cfg(test)]
mod test {
    use template::{Template};
    use registry::{Registry};
    use render::{RenderContext, Renderable, RenderError, Helper};
    use helpers::{HelperDef};
    use context::{Context};
    use support::str::StringWriter;

    #[derive(Clone, Copy)]
    struct DummyHelper;

    impl HelperDef for DummyHelper {
        fn call(&self, c: &Context, h: &Helper, r: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
            try!(h.template().unwrap().render(c, r, rc));
            Ok(())
        }
    }

    static DUMMY_HELPER: DummyHelper = DummyHelper;

    #[test]
    fn test_registry_operations() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let t2 = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index2", t2.clone());

        assert_eq!((*r.get_template(&("index".to_string())).unwrap()).to_string(),
                   t.to_string());
        assert_eq!(r.templates.len(), 2);

        r.unregister_template(&("index".to_string()));
        assert_eq!(r.templates.len(), 1);

        r.clear_templates();
        assert_eq!(r.templates.len(), 0);

        r.register_helper("dummy", Box::new(DUMMY_HELPER));

        // built-in helpers plus 1
        assert_eq!(r.helpers.len(), 10+1);
    }

    #[test]
    fn test_renderw() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.renderw("index", &context, &mut sw).ok().unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

    }

    #[test]
    fn test_escape_fn() {
        let mut r = Registry::new();

        let input = String::from("\"<>&");

        r.register_template_string("test", String::from("{{this}}")).unwrap();

        assert_eq!(""<>&", r.render("test", &input).unwrap());

        r.register_escape_fn(|s| s.into());

        assert_eq!("\"<>&", r.render("test", &input).unwrap());

        r.unregister_escape_fn();

        assert_eq!(""<>&", r.render("test", &input).unwrap());
    }

    #[test]
    fn test_template_render() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        assert_eq!("

".to_string(),
                   r.template_render("{{> index}}", &{}).unwrap());

        assert_eq!("hello world".to_string(),
                   r.template_render("hello {{this}}", &"world".to_string()).unwrap());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.template_renderw("{{> index}}", &context, &mut sw).unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());
    }
}
<commit_msg>(test) added test for error cases<commit_after>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;

#[cfg(not(feature = "serde_type"))]
use serialize::json::ToJson;
#[cfg(feature = "serde_type")]
use serde::ser::Serialize as ToJson;

use template::{Template};
use render::{Renderable, RenderError, RenderContext};
use helpers::{HelperDef};
use context::{Context};
use helpers;
use support::str::StringWriter;
use error::{TemplateError, TemplateFileError, TemplateRenderError};

/// This type represents an *escape fn*, that is a function who's purpose it is
/// to escape potentially problematic characters in a string.
///
/// The default *escape fn* replaces the characters `&"<>`
/// with the equivalent html / xml entities.
///
/// An *escape fn* is represented as a `Box` to avoid unnecessary type
/// parameters (and because traits cannot be aliased using `type`).
pub type EscapeFn = Box<Fn(&str) -> String + Send + Sync>;

fn get_default_escape_fn() -> EscapeFn {
    Box::new(|data| {
        data.replace("&", "&")
            .replace("\"", """)
            .replace("<", "<")
            .replace(">", ">")
    })
}

/// The single entry point of your Handlebars templates
///
/// It maintains compiled templates and registered helpers.
pub struct Registry {
    templates: HashMap<String, Template>,
    helpers: HashMap<String, Box<HelperDef + 'static>>,
    escape_fn: EscapeFn
}

impl Registry {
    pub fn new() -> Registry {
        let mut r = Registry {
            templates: HashMap::new(),
            helpers: HashMap::new(),
            escape_fn: get_default_escape_fn(),
        };

        r.register_helper("if", Box::new(helpers::IF_HELPER));
        r.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        r.register_helper("each", Box::new(helpers::EACH_HELPER));
        r.register_helper("with", Box::new(helpers::WITH_HELPER));
        r.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        r.register_helper("raw", Box::new(helpers::RAW_HELPER));
        r.register_helper(">", Box::new(helpers::INCLUDE_HELPER));
        r.register_helper("block", Box::new(helpers::BLOCK_HELPER));
        r.register_helper("partial", Box::new(helpers::PARTIAL_HELPER));
        r.register_helper("log", Box::new(helpers::LOG_HELPER));

        r
    }

    /// Register a template
    pub fn register_template(&mut self, name: &str, mut template: Template) {
        template.name = Some(name.to_owned());
        self.templates.insert(name.to_string(), template);
    }

    /// Register a template string
    pub fn register_template_string(&mut self, name: &str, tpl_str: String) -> Result<(), TemplateError> {
        try!(Template::compile_with_name(tpl_str, name.to_owned())
             .and_then(|t| Ok(self.templates.insert(name.to_string(), t))));
        Ok(())
    }

    /// Register a template from a path
    pub fn register_template_file(&mut self, name: &str, tpl_path: &Path) -> Result<(), TemplateFileError> {
        let mut file = try!(File::open(tpl_path));
        self.register_template_source(name, &mut file)
    }

    /// Register a template from `std::io::Read` source
    pub fn register_template_source(&mut self, name: &str, tpl_source: &mut Read) -> Result<(), TemplateFileError> {
        let mut buf = String::new();
        try!(tpl_source.read_to_string(&mut buf));
        try!(Template::compile_with_name(&buf, name.to_owned())
             .and_then(|t| Ok(self.templates.insert(name.to_string(), t))));
        Ok(())
    }

    /// remove a template from the registry
    pub fn unregister_template(&mut self, name: &str) {
        self.templates.remove(name);
    }

    /// register a helper
    pub fn register_helper(&mut self, name: &str, def: Box<HelperDef + 'static>) -> Option<Box<HelperDef + 'static>> {
        self.helpers.insert(name.to_string(), def)
    }

    /// Register a new *escape fn* to be used from now on by this registry.
    pub fn register_escape_fn<F: 'static + Fn(&str) -> String + Send + Sync>(&mut self, escape_fn: F) {
        self.escape_fn = Box::new(escape_fn);
    }

    /// Restore the default *escape fn*.
    pub fn unregister_escape_fn(&mut self) {
        self.escape_fn = get_default_escape_fn();
    }

    /// Get a reference to the current *escape fn*.
    pub fn get_escape_fn(&self) -> &Fn(&str) -> String {
        &*self.escape_fn
    }

    /// Return a registered template,
    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    /// Return a registered helper
    pub fn get_helper(&self, name: &str) -> Option<&Box<HelperDef + 'static>> {
        self.helpers.get(name)
    }

    /// Return all templates registered
    pub fn get_templates(&self) -> &HashMap<String, Template> {
        &self.templates
    }

    /// Unregister all templates
    pub fn clear_templates(&mut self) {
        self.templates.clear();
    }

    /// Render a registered template with some data into a string
    pub fn render<T>(&self, name: &str, ctx: &T) -> Result<String, RenderError> where T: ToJson {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.renderw(name, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// Render a registered template with some data to the `std::io::Write`
    pub fn renderw(&self, name: &str, context: &Context, writer: &mut Write) -> Result<(), RenderError> {
        let template = self.get_template(&name.to_string());

        if let Some(t) = template {
            let mut render_context = RenderContext::new(writer);
            render_context.root_template = t.name.clone();
            (*t).render(context, self, &mut render_context)
        } else {
            Err(RenderError::new(format!("Template not found: {}", name)))
        }
    }

    /// render a template string using current registry without register it
    pub fn template_render<T>(&self, template_string: &str, ctx: &T) -> Result<String, TemplateRenderError> where T: ToJson {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.template_renderw(template_string, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// render a template string using current registry without register it
    pub fn template_renderw(&self, template_string: &str, context: &Context, writer: &mut Write) -> Result<(), TemplateRenderError> {
        let tpl = try!(Template::compile(template_string).map_err(TemplateRenderError::from));
        let mut render_context = RenderContext::new(writer);
        tpl.render(context, self, &mut render_context).map_err(TemplateRenderError::from)
    }
}

#[cfg(test)]
mod test {
    use template::{Template};
    use registry::{Registry};
    use render::{RenderContext, Renderable, RenderError, Helper};
    use helpers::{HelperDef};
    use context::{Context};
    use support::str::StringWriter;
    use error::TemplateRenderError;

    #[derive(Clone, Copy)]
    struct DummyHelper;

    impl HelperDef for DummyHelper {
        fn call(&self, c: &Context, h: &Helper, r: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
            try!(h.template().unwrap().render(c, r, rc));
            Ok(())
        }
    }

    static DUMMY_HELPER: DummyHelper = DummyHelper;

    #[test]
    fn test_registry_operations() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let t2 = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index2", t2.clone());

        assert_eq!((*r.get_template(&("index".to_string())).unwrap()).to_string(),
                   t.to_string());
        assert_eq!(r.templates.len(), 2);

        r.unregister_template(&("index".to_string()));
        assert_eq!(r.templates.len(), 1);

        r.clear_templates();
        assert_eq!(r.templates.len(), 0);

        r.register_helper("dummy", Box::new(DUMMY_HELPER));

        // built-in helpers plus 1
        assert_eq!(r.helpers.len(), 10+1);
    }

    #[test]
    fn test_renderw() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.renderw("index", &context, &mut sw).ok().unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

    }

    #[test]
    fn test_escape_fn() {
        let mut r = Registry::new();

        let input = String::from("\"<>&");

        r.register_template_string("test", String::from("{{this}}")).unwrap();

        assert_eq!(""<>&", r.render("test", &input).unwrap());

        r.register_escape_fn(|s| s.into());

        assert_eq!("\"<>&", r.render("test", &input).unwrap());

        r.unregister_escape_fn();

        assert_eq!(""<>&", r.render("test", &input).unwrap());
    }

    #[test]
    fn test_template_render() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        assert_eq!("

".to_string(),
                   r.template_render("{{> index}}", &{}).unwrap());

        assert_eq!("hello world".to_string(),
                   r.template_render("hello {{this}}", &"world".to_string()).unwrap());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.template_renderw("{{> index}}", &context, &mut sw).unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

        // fail for template error
        match r.template_render("{{ hello", &{}).unwrap_err() {
            TemplateRenderError::TemplateError(_) => {},
            _ => { panic!(); }
        }

        // fail to render error
        match r.template_render("{{> notfound}}", &{}).unwrap_err() {
            TemplateRenderError::RenderError(_) => {},
            _ => { panic!(); }
        }
    }
}
<|endoftext|>


<commit_before>use std::sync::{Arc, Weak};
use std::default::Default;

use mioco;
use mioco::sync::Mutex;
use mioco::sync::mpsc::{channel, Receiver};
use trust_dns::op::{Message, ResponseCode, Query};
use trust_dns::rr::{Name};
use itertools::Itertools;

use utils::{Result, MessageExt, AsDisplay, Future, CloneExt};
use cache::Cache;
use nscache::NsCache;
use config::Config;
use recursive::RecursiveResolver;
use forwarding::ForwardingResolver;
use query::query_multiple_handle_futures;
use control::{ControlServer, Error as ControlError, ControlResult};

// Idea: Use DNSSEC to check whether a domain is poisoned by GFW

#[derive(Debug, Clone)]
pub enum ErrorKind {
    NoNameserver,
    EmptyAnswer,
    InsaneNsReferral,
    /// Another coroutine is already working on the query
    LostRace,
    RejectedIp,
    AlreadyQueried,
}
pub struct Resolver {
    pub cache: Cache,
    pub ns_cache: NsCache,
    pub config: Config,
    forwarders: Vec<Arc<ForwardingResolver>>,
    control_server: Mutex<ControlServer>,
}
custom_derive! {
    #[derive(Clone, NewtypeFrom, NewtypeDeref)]
    pub struct RcResolver(Arc<Resolver>);
}
custom_derive! {
    #[derive(Clone, NewtypeFrom)]
    pub struct RcResolverWeak(Weak<Resolver>);
}

impl RcResolver {
    pub fn new(config: Config) -> Self {
        let forwarders = ForwardingResolver::create_all(&config);
        let (send, recv) = channel::<Query>();
        let ret: RcResolver = Arc::new(Resolver {
            cache: Cache::with_expiration_notifier(send),
            ns_cache: NsCache::default(),
            config: config,
            forwarders: forwarders,
            control_server: Mutex::new(ControlServer::new()),
        }).into();
        ret.control_server.lock().unwrap().run(&ret)
        .expect("Failed to initialize control socket");
        ret.init_root_servers();
        ret.attach();
        ret.handle_cache_expiration_channel(recv);
        ret
    }
    pub fn to_weak(&self) -> RcResolverWeak {
        RcResolverWeak(Arc::downgrade(self))
    }
    pub fn current_weak() -> RcResolverWeak {
        mioco::get_userdata::<RcResolverWeak>()
        .map_or_else(RcResolverWeak::empty, |x| (*x).clone())
    }
    pub fn current() -> Self {
        Self::current_weak().upgrade().expect("Called `current()` without initialization")
    }
    fn attach(&self) {
        assert!(RcResolver::current_weak().upgrade().is_none());
        mioco::set_userdata(self.to_weak());
        mioco::set_children_userdata(Some(self.to_weak()));
    }
}
impl Default for RcResolver {
    fn default() -> Self {
        Self::new(Config::default())
    }
}
impl RcResolverWeak {
    pub fn upgrade(&self) -> Option<RcResolver> {
        self.0.upgrade().map(|x| x.into())
    }
    fn empty() -> Self {
        RcResolverWeak(Weak::new())
    }
}
impl RcResolver {
    fn init_root_servers(&self) {
        let mut guard = self.ns_cache.lock().unwrap();
        let mut entry = guard.lookup_or_insert(&Name::root());
        entry.pin();
        for ip in &self.config.root_servers {
            entry.add_ns(*ip, Name::root(), None);
        }
    }
    pub fn handle_control_command(&self, cmd: &str, _: &[String]) -> ControlResult {
        match cmd {
            "ping" => Ok("Pong".into()),
            "expire-cache" => {
                self.cache.expire_all();
                Ok("All cache entries have been marked as expired".into())
            },
            "clear-cache" => {
                self.cache.clear();
                Ok("Cleared cache".into())
            },
            _ => Err(ControlError::UnknownCommand),
        }
    }
    fn handle_cache_expiration_channel(&self, ch: Receiver<Query>) {
        let resolver_weak = self.to_weak();
        mioco::spawn(move || {
            while let Ok(q) = ch.recv() {
                if let Some(res) = resolver_weak.upgrade() {
                    let updated = match RecursiveResolver::resolve(&q, res.clone().into()) {
                        Ok(msg) => [ResponseCode::NoError, ResponseCode::NXDomain]
                                   .contains(&msg.get_response_code()),
                        Err(_) => false,
                    };
                    if !updated {
                        res.cache.operate(q.get_name(), |entry| {
                            let still_expired = entry.lookup_entry(q.get_query_type())
                            .map_or(false, |x| x.is_expired());
                            if still_expired {
                                debug!("Purging cache entry due to failed update: {}",
                                       q.as_disp());
                                entry.purge(q.get_query_type());
                            }
                        });
                    }
                } else {
                    break;
                }
            }
            debug!("Resolver is dropped, cache update coroutine is exiting");
        });
    }
    fn resolve_recursive(self, q: Query) -> Result<Message> {
        RecursiveResolver::resolve(&q, self)
    }
    fn resolve_internal(&self, q: &Query) -> Result<Message> {
        let mut futures = self.forwarders.iter().cloned().map(move |forwarder| {
            let qc = (*q).clone();
            let res = self.clone();
            Future::from_fn(move || forwarder.resolve(qc, res))
        }).collect_vec();
        let qc = (*q).clone();
        let res = self.clone();
        futures.push(Future::from_fn(move || res.resolve_recursive(qc)));
        query_multiple_handle_futures(&mut futures)
    }
    pub fn resolve(&self, msg: &mut Message) -> Result<()> {
        debug_assert!(msg.get_queries().len() == 1);
        debug_assert!(msg.get_answers().is_empty());
        debug_assert!(msg.get_name_servers().is_empty());
        let query_type = msg.get_queries()[0].get_query_type();
        let name = msg.get_queries()[0].get_name().clone();
        let cache_hit = self.cache.lookup_with_type(
            &name,
            query_type,
            |cached| msg.copy_resp_from(&cached),
        ).is_some();
        if cache_hit {
            return Ok(());
        }
        let resp = try!(self.resolve_internal(&msg.get_queries()[0]));
        msg.copy_resp_from(&resp);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use trust_dns::op::*;
    use trust_dns::rr::*;
    use mioco;

    use super::*;
    use ::mioco_config_start;
    use ::config::*;

    #[test]
    fn simple_recursive_query() {
        mioco_config_start(|| {
            let config = Config::default();
            let resolver = RcResolver::new(config);
            let mut q = Query::new();
            q.name(Name::parse("www.google.com", Some(&Name::root())).unwrap());
            let result = resolver.resolve_recursive(q).unwrap();
            assert!(result.get_answers().len() > 0);
        }).unwrap();
    }
    #[test]
    #[allow(unused_variables)]
    fn current() {
        mioco_config_start(|| {
            let config = Config::default();
            let resolver = RcResolver::new(config);
            RcResolver::current();
            mioco::spawn(move || {
                RcResolver::current();
                mioco::spawn(move || {
                    RcResolver::current();
                }).join().unwrap();
            }).join().unwrap();
        }).unwrap();
    }
    #[test]
    #[should_panic]
    #[allow(unused_variables)]
    fn no_multiple_instances() {
        // FIXME: We have to use mioco::start here, using `mioco_config_start` (which
        // doesn't handle panics in coroutines) will crash the test process
        mioco::start(|| {
            let config = Config::default();
            let resolver = RcResolver::new(config);
            let config = Config::default();
            let resolver2 = RcResolver::new(config);
        }).unwrap();
    }
}
<commit_msg>Don't run control server when running tests<commit_after>use std::sync::{Arc, Weak};
use std::default::Default;

use mioco;
use mioco::sync::Mutex;
use mioco::sync::mpsc::{channel, Receiver};
use trust_dns::op::{Message, ResponseCode, Query};
use trust_dns::rr::{Name};
use itertools::Itertools;

use utils::{Result, MessageExt, AsDisplay, Future, CloneExt};
use cache::Cache;
use nscache::NsCache;
use config::Config;
use recursive::RecursiveResolver;
use forwarding::ForwardingResolver;
use query::query_multiple_handle_futures;
use control::{ControlServer, Error as ControlError, ControlResult};

// Idea: Use DNSSEC to check whether a domain is poisoned by GFW

#[derive(Debug, Clone)]
pub enum ErrorKind {
    NoNameserver,
    EmptyAnswer,
    InsaneNsReferral,
    /// Another coroutine is already working on the query
    LostRace,
    RejectedIp,
    AlreadyQueried,
}
pub struct Resolver {
    pub cache: Cache,
    pub ns_cache: NsCache,
    pub config: Config,
    forwarders: Vec<Arc<ForwardingResolver>>,
    control_server: Mutex<ControlServer>,
}
custom_derive! {
    #[derive(Clone, NewtypeFrom, NewtypeDeref)]
    pub struct RcResolver(Arc<Resolver>);
}
custom_derive! {
    #[derive(Clone, NewtypeFrom)]
    pub struct RcResolverWeak(Weak<Resolver>);
}

impl RcResolver {
    pub fn new(config: Config) -> Self {
        let forwarders = ForwardingResolver::create_all(&config);
        let (send, recv) = channel::<Query>();
        let ret: RcResolver = Arc::new(Resolver {
            cache: Cache::with_expiration_notifier(send),
            ns_cache: NsCache::default(),
            config: config,
            forwarders: forwarders,
            control_server: Mutex::new(ControlServer::new()),
        }).into();
        if cfg!(not(test)) {
            ret.control_server.lock().unwrap().run(&ret)
            .expect("Failed to initialize control socket");
        }
        ret.init_root_servers();
        ret.attach();
        ret.handle_cache_expiration_channel(recv);
        ret
    }
    #[cfg(test)]
    pub fn run_control_server(&self) {
        self.control_server.lock().unwrap().run(self)
        .expect("Failed to initialize control socket");
    }
    pub fn to_weak(&self) -> RcResolverWeak {
        RcResolverWeak(Arc::downgrade(self))
    }
    pub fn current_weak() -> RcResolverWeak {
        mioco::get_userdata::<RcResolverWeak>()
        .map_or_else(RcResolverWeak::empty, |x| (*x).clone())
    }
    pub fn current() -> Self {
        Self::current_weak().upgrade().expect("Called `current()` without initialization")
    }
    fn attach(&self) {
        assert!(RcResolver::current_weak().upgrade().is_none());
        mioco::set_userdata(self.to_weak());
        mioco::set_children_userdata(Some(self.to_weak()));
    }
}
impl Default for RcResolver {
    fn default() -> Self {
        Self::new(Config::default())
    }
}
impl RcResolverWeak {
    pub fn upgrade(&self) -> Option<RcResolver> {
        self.0.upgrade().map(|x| x.into())
    }
    fn empty() -> Self {
        RcResolverWeak(Weak::new())
    }
}
impl RcResolver {
    fn init_root_servers(&self) {
        let mut guard = self.ns_cache.lock().unwrap();
        let mut entry = guard.lookup_or_insert(&Name::root());
        entry.pin();
        for ip in &self.config.root_servers {
            entry.add_ns(*ip, Name::root(), None);
        }
    }
    pub fn handle_control_command(&self, cmd: &str, _: &[String]) -> ControlResult {
        match cmd {
            "ping" => Ok("Pong".into()),
            "expire-cache" => {
                self.cache.expire_all();
                Ok("All cache entries have been marked as expired".into())
            },
            "clear-cache" => {
                self.cache.clear();
                Ok("Cleared cache".into())
            },
            _ => Err(ControlError::UnknownCommand),
        }
    }
    fn handle_cache_expiration_channel(&self, ch: Receiver<Query>) {
        let resolver_weak = self.to_weak();
        mioco::spawn(move || {
            while let Ok(q) = ch.recv() {
                if let Some(res) = resolver_weak.upgrade() {
                    let updated = match RecursiveResolver::resolve(&q, res.clone().into()) {
                        Ok(msg) => [ResponseCode::NoError, ResponseCode::NXDomain]
                                   .contains(&msg.get_response_code()),
                        Err(_) => false,
                    };
                    if !updated {
                        res.cache.operate(q.get_name(), |entry| {
                            let still_expired = entry.lookup_entry(q.get_query_type())
                            .map_or(false, |x| x.is_expired());
                            if still_expired {
                                debug!("Purging cache entry due to failed update: {}",
                                       q.as_disp());
                                entry.purge(q.get_query_type());
                            }
                        });
                    }
                } else {
                    break;
                }
            }
            debug!("Resolver is dropped, cache update coroutine is exiting");
        });
    }
    fn resolve_recursive(self, q: Query) -> Result<Message> {
        RecursiveResolver::resolve(&q, self)
    }
    fn resolve_internal(&self, q: &Query) -> Result<Message> {
        let mut futures = self.forwarders.iter().cloned().map(move |forwarder| {
            let qc = (*q).clone();
            let res = self.clone();
            Future::from_fn(move || forwarder.resolve(qc, res))
        }).collect_vec();
        let qc = (*q).clone();
        let res = self.clone();
        futures.push(Future::from_fn(move || res.resolve_recursive(qc)));
        query_multiple_handle_futures(&mut futures)
    }
    pub fn resolve(&self, msg: &mut Message) -> Result<()> {
        debug_assert!(msg.get_queries().len() == 1);
        debug_assert!(msg.get_answers().is_empty());
        debug_assert!(msg.get_name_servers().is_empty());
        let query_type = msg.get_queries()[0].get_query_type();
        let name = msg.get_queries()[0].get_name().clone();
        let cache_hit = self.cache.lookup_with_type(
            &name,
            query_type,
            |cached| msg.copy_resp_from(&cached),
        ).is_some();
        if cache_hit {
            return Ok(());
        }
        let resp = try!(self.resolve_internal(&msg.get_queries()[0]));
        msg.copy_resp_from(&resp);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use trust_dns::op::*;
    use trust_dns::rr::*;
    use mioco;

    use super::*;
    use ::mioco_config_start;
    use ::config::*;

    #[test]
    fn simple_recursive_query() {
        mioco_config_start(|| {
            let config = Config::default();
            let resolver = RcResolver::new(config);
            let mut q = Query::new();
            q.name(Name::parse("www.google.com", Some(&Name::root())).unwrap());
            let result = resolver.resolve_recursive(q).unwrap();
            assert!(result.get_answers().len() > 0);
        }).unwrap();
    }
    #[test]
    #[allow(unused_variables)]
    fn current() {
        mioco_config_start(|| {
            let config = Config::default();
            let resolver = RcResolver::new(config);
            RcResolver::current();
            mioco::spawn(move || {
                RcResolver::current();
                mioco::spawn(move || {
                    RcResolver::current();
                }).join().unwrap();
            }).join().unwrap();
        }).unwrap();
    }
    #[test]
    #[should_panic]
    #[allow(unused_variables)]
    fn no_multiple_instances() {
        // FIXME: We have to use mioco::start here, using `mioco_config_start` (which
        // doesn't handle panics in coroutines) will crash the test process
        mioco::start(|| {
            let config = Config::default();
            let resolver = RcResolver::new(config);
            let config = Config::default();
            let resolver2 = RcResolver::new(config);
        }).unwrap();
    }
}
<|endoftext|>


<commit_before>// FIXME: add the support of transient objects

use any_cache::{Cache, HashCache};
use notify::{self, RecommendedWatcher, Watcher};
use std::any::Any;
use std::cell::RefCell;
use std::collections::HashMap;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::channel;
use std::thread;
use std::time::{Duration, Instant};

/// Class of types that can be loaded.
pub trait Load: Sized {
  /// Arguments passed at loading.
  type Args: Clone;

  /// Static string representing the name of the type (used in filesystem caching).
  const TY_STR: &'static str;

  fn load<P>(path: P, cache: &mut ResCache, args: Self::Args) -> Result<Self> where P: AsRef<Path>;
}

/// Class of types that can be reloaded.
///
/// The idea is to simply recover the arguments used in `Load::load`.
pub trait Reload: Load {
  fn reload_args(&self) -> Self::Args;
}

/// Default implementation for types which are loaded without any arguments.
impl<T> Reload for T where T: Load<Args=()> {
  fn reload_args(&self) -> Self::Args {
    ()
  }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LoadError {
  FileNotFound(PathBuf, String),
  ParseFailed(String),
  ConversionFailed(String)
}

pub type Result<T> = ::std::result::Result<T, LoadError>;

/// Resources are wrapped in this type.
#[derive(Debug)]
pub struct Res<T>(Rc<RefCell<T>>);

impl<T> Res<T> {
  pub fn as_rc(&self) -> &Rc<RefCell<T>> {
    &self.0
  }
}

impl<T> Clone for Res<T> {
  fn clone(&self) -> Self {
    Res(self.0.clone())
  }
}

impl<T> Deref for Res<T> {
  type Target = RefCell<T>;

  fn deref(&self) -> &Self::Target {
    &self.0
  }
}

/// Time to await after a resource update to establish that it should be reloaded.
const UPDATE_AWAIT_TIME_MS: u64 = 1000;

/// Resource cache. Responsible for caching resource.
pub struct ResCache {
  // contains all the typed-erased Rc<RefCell<T>>
  cache: HashCache<PathBuf>,
  // contains all metadata on resources
  metadata: HashMap<PathBuf, ResMetaData>,
  // vector of pair (path, timestamp) giving indication on resources to reload
  dirty: Arc<Mutex<Vec<(PathBuf, Instant)>>>,
  #[allow(dead_code)]
  watcher_thread: thread::JoinHandle<()>
}

struct ResMetaData {
  on_reload: Box<Fn(&mut ResCache)>,
  last_update_instant: Instant
}

impl ResCache {
  /// Create a new cache.
  pub fn new<P>(root: P) -> Self where P: AsRef<Path> {
    let dirty: Arc<Mutex<Vec<(PathBuf, Instant)>>> = Arc::new(Mutex::new(Vec::new()));
    let dirty_ = dirty.clone();

    let root = root.as_ref().to_owned();
    let (wsx, wrx) = channel();
    let mut watcher: RecommendedWatcher = Watcher::new(wsx).unwrap();

    let join_handle = thread::spawn(move || {
      let _ = watcher.watch(root);

      for event in wrx.iter() {
        if let notify::Event { path: Some(path), op: Ok(notify::op::WRITE) } = event {
          dirty_.lock().unwrap().push((path.clone(), Instant::now()));
        }
      }
    });

    ResCache {
      cache: HashCache::new(),
      metadata: HashMap::new(),
      dirty: dirty,
      watcher_thread: join_handle
    }
  }

  /// Get a resource from the cache.
  pub fn get<T>(&mut self, key: &str, args: T::Args) -> Option<Res<T>> where T: 'static + Any + Reload {
    let path_str = format!("data/{}/{}", T::TY_STR, key);
    let path = Path::new(&path_str);
    let path_buf = path.to_owned();

    match self.cache.get::<Res<T>>(&path_buf).cloned() {
      r@Some(..) => {
        deb!("cache hit for {} ({})", key, path_str);
        r
      },
      None => {
        deb!("cache miss for {} ({})", key, path_str);

        // specific loading
        if path.exists() {
          match T::load(&path, self, args.clone()) {
            Ok(resource) => {
              let res = Res(Rc::new(RefCell::new(resource)));
              let res_ = res.clone();

              let path_buf_ = path_buf.clone();
              // closure used to reload the object when needed
              let on_reload: Box<for<'a> Fn(&'a mut ResCache)> = Box::new(move |cache_| {
                match T::load(&path_buf_, cache_, args.clone()) {
                  Ok(new_resource) => {
                    // replace the current resource with the freshly loaded one
                    *res_.borrow_mut() = new_resource;
                    deb!("reloaded resource from {:?}", path_buf_);
                  },
                  Err(e) => {
                    warn!("reloading resource from {:?} has failed:\n{:#?}", path_buf_, e);
                  }
                }
              });

              let metadata = ResMetaData {
                on_reload: on_reload,
                last_update_instant: Instant::now()
              };

              // cache the resource and its meta data
              self.cache.save(path_buf.clone(), res.clone());
              self.metadata.insert(path_buf, metadata);

              Some(res)
            },
            Err(e) => {
              err!("unable to load resource from {}:\n{:#?}", path_str, e);
              None
            }
          }
        } else { // path doesn’t exist
          err!("resource at {} cannot be found", path_str);
          None
        }
      }
    }
  }

  /// Synchronize the cache by updating the resource that ought to.
  pub fn sync(&mut self) {
    let dirty = self.dirty.clone();
    let mut dirty_ = dirty.lock().unwrap();

    for &(ref path, ref instant) in dirty_.iter() {
      let mut metadata = self.metadata.remove(path).unwrap();

      if instant.duration_since(metadata.last_update_instant) >= Duration::from_millis(UPDATE_AWAIT_TIME_MS) {
        (metadata.on_reload)(self);
      }

      metadata.last_update_instant = *instant;
      self.metadata.insert(path.clone(), metadata);
    }

    dirty_.clear();
  }
}
<commit_msg>From<Rc…> for Res…<commit_after>// FIXME: add the support of transient objects

use any_cache::{Cache, HashCache};
use notify::{self, RecommendedWatcher, Watcher};
use std::any::Any;
use std::cell::RefCell;
use std::collections::HashMap;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::channel;
use std::thread;
use std::time::{Duration, Instant};

/// Class of types that can be loaded.
pub trait Load: Sized {
  /// Arguments passed at loading.
  type Args: Clone;

  /// Static string representing the name of the type (used in filesystem caching).
  const TY_STR: &'static str;

  fn load<P>(path: P, cache: &mut ResCache, args: Self::Args) -> Result<Self> where P: AsRef<Path>;
}

/// Class of types that can be reloaded.
///
/// The idea is to simply recover the arguments used in `Load::load`.
pub trait Reload: Load {
  fn reload_args(&self) -> Self::Args;
}

/// Default implementation for types which are loaded without any arguments.
impl<T> Reload for T where T: Load<Args=()> {
  fn reload_args(&self) -> Self::Args {
    ()
  }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LoadError {
  FileNotFound(PathBuf, String),
  ParseFailed(String),
  ConversionFailed(String)
}

pub type Result<T> = ::std::result::Result<T, LoadError>;

/// Resources are wrapped in this type.
#[derive(Debug)]
pub struct Res<T>(Rc<RefCell<T>>);

impl<T> Res<T> {
  pub fn as_rc(&self) -> &Rc<RefCell<T>> {
    &self.0
  }
}

impl<T> Clone for Res<T> {
  fn clone(&self) -> Self {
    Res(self.0.clone())
  }
}

impl<T> Deref for Res<T> {
  type Target = RefCell<T>;

  fn deref(&self) -> &Self::Target {
    &self.0
  }
}

impl<T> From<Rc<RefCell<T>>> for Res<T> {
  fn from(rc: Rc<RefCell<T>>) -> Self {
    Res(rc)
  }
}

/// Time to await after a resource update to establish that it should be reloaded.
const UPDATE_AWAIT_TIME_MS: u64 = 1000;

/// Resource cache. Responsible for caching resource.
pub struct ResCache {
  // contains all the typed-erased Rc<RefCell<T>>
  cache: HashCache<PathBuf>,
  // contains all metadata on resources
  metadata: HashMap<PathBuf, ResMetaData>,
  // vector of pair (path, timestamp) giving indication on resources to reload
  dirty: Arc<Mutex<Vec<(PathBuf, Instant)>>>,
  #[allow(dead_code)]
  watcher_thread: thread::JoinHandle<()>
}

struct ResMetaData {
  on_reload: Box<Fn(&mut ResCache)>,
  last_update_instant: Instant
}

impl ResCache {
  /// Create a new cache.
  pub fn new<P>(root: P) -> Self where P: AsRef<Path> {
    let dirty: Arc<Mutex<Vec<(PathBuf, Instant)>>> = Arc::new(Mutex::new(Vec::new()));
    let dirty_ = dirty.clone();

    let root = root.as_ref().to_owned();
    let (wsx, wrx) = channel();
    let mut watcher: RecommendedWatcher = Watcher::new(wsx).unwrap();

    let join_handle = thread::spawn(move || {
      let _ = watcher.watch(root);

      for event in wrx.iter() {
        if let notify::Event { path: Some(path), op: Ok(notify::op::WRITE) } = event {
          dirty_.lock().unwrap().push((path.clone(), Instant::now()));
        }
      }
    });

    ResCache {
      cache: HashCache::new(),
      metadata: HashMap::new(),
      dirty: dirty,
      watcher_thread: join_handle
    }
  }

  /// Get a resource from the cache.
  pub fn get<T>(&mut self, key: &str, args: T::Args) -> Option<Res<T>> where T: 'static + Any + Reload {
    let path_str = format!("data/{}/{}", T::TY_STR, key);
    let path = Path::new(&path_str);
    let path_buf = path.to_owned();

    match self.cache.get::<Res<T>>(&path_buf).cloned() {
      r@Some(..) => {
        deb!("cache hit for {} ({})", key, path_str);
        r
      },
      None => {
        deb!("cache miss for {} ({})", key, path_str);

        // specific loading
        if path.exists() {
          match T::load(&path, self, args.clone()) {
            Ok(resource) => {
              let res = Res(Rc::new(RefCell::new(resource)));
              let res_ = res.clone();

              let path_buf_ = path_buf.clone();
              // closure used to reload the object when needed
              let on_reload: Box<for<'a> Fn(&'a mut ResCache)> = Box::new(move |cache_| {
                match T::load(&path_buf_, cache_, args.clone()) {
                  Ok(new_resource) => {
                    // replace the current resource with the freshly loaded one
                    *res_.borrow_mut() = new_resource;
                    deb!("reloaded resource from {:?}", path_buf_);
                  },
                  Err(e) => {
                    warn!("reloading resource from {:?} has failed:\n{:#?}", path_buf_, e);
                  }
                }
              });

              let metadata = ResMetaData {
                on_reload: on_reload,
                last_update_instant: Instant::now()
              };

              // cache the resource and its meta data
              self.cache.save(path_buf.clone(), res.clone());
              self.metadata.insert(path_buf, metadata);

              Some(res)
            },
            Err(e) => {
              err!("unable to load resource from {}:\n{:#?}", path_str, e);
              None
            }
          }
        } else { // path doesn’t exist
          err!("resource at {} cannot be found", path_str);
          None
        }
      }
    }
  }

  /// Synchronize the cache by updating the resource that ought to.
  pub fn sync(&mut self) {
    let dirty = self.dirty.clone();
    let mut dirty_ = dirty.lock().unwrap();

    for &(ref path, ref instant) in dirty_.iter() {
      let mut metadata = self.metadata.remove(path).unwrap();

      if instant.duration_since(metadata.last_update_instant) >= Duration::from_millis(UPDATE_AWAIT_TIME_MS) {
        (metadata.on_reload)(self);
      }

      metadata.last_update_instant = *instant;
      self.metadata.insert(path.clone(), metadata);
    }

    dirty_.clear();
  }
}
<|endoftext|>