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

use serde::Serialize;

use regex::{Captures, Regex};

use template::Template;
use render::{RenderContext, Renderable};
use context::Context;
use helpers::{self, HelperDef};
use directives::{self, DirectiveDef};
use support::str::StringWriter;
use error::{RenderError, 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.get(0).map(|m| m.as_str()) {
                Some("<") => "<",
                Some(">") => ">",
                Some("\"") => """,
                Some("&") => "&",
                _ => unreachable!(),
            }.to_owned()
        })
        .into_owned()
}

/// `EscapeFn` that do not 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()
    }

    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("log", Box::new(helpers::LOG_HELPER));

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

    /// Enable handlebars template source map
    ///
    /// Source map provides line/col reporting on error. It uses slightly
    /// more memory to maintain the data.
    ///
    /// Default is true.
    pub fn source_map_enable(&mut self, enable: bool) {
        self.source_map = enable;
    }

    /// Register a template string
    ///
    /// Returns `TemplateError` if there is syntax error on parsing template.
    pub fn register_template_string<S>(
        &mut self,
        name: &str,
        tpl_str: S,
    ) -> Result<(), TemplateError>
    where
        S: AsRef<str>,
    {
        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
    ///
    /// A named partial will be added to the registry. It will overwrite template with
    /// same name. Currently registered partial is just identical to template.
    pub fn register_partial<S>(&mut self, name: &str, partial_str: S) -> Result<(), TemplateError>
    where
        S: AsRef<str>,
    {
        self.register_template_string(name, partial_str)
    }

    /// Register a template from a path
    pub fn register_template_file
(
        &mut self,
        name: &str,
        tpl_path: P,
    ) -> Result<(), TemplateFileError>
    where
        P: AsRef<Path>,
    {
        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 decorator
    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
    ///
    /// * `name` is the template name you registred previously
    /// * `ctx` is the data that implements `serde::Serialize`
    ///
    /// Returns rendered string or an struct with error information
    pub fn render<T>(&self, name: &str, data: &T) -> Result<String, RenderError>
    where
        T: Serialize,
    {
        let mut writer = StringWriter::new();
        {
            try!(self.renderw(name, data, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// Render a registered template and write some data to the `std::io::Write`
    pub fn renderw<T>(&self, name: &str, data: &T, writer: &mut Write) -> Result<(), RenderError>
    where
        T: Serialize,
    {
        self.get_template(&name.to_string())
            .ok_or(RenderError::new(format!("Template not found: {}", name)))
            .and_then(|t| {
                let ctx = try!(Context::wraps(data));
                let mut local_helpers = HashMap::new();
                let mut render_context = RenderContext::new(ctx, &mut local_helpers, writer);
                render_context.root_template = t.name.clone();
                t.render(self, &mut render_context)
            })
    }

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

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

    /// render a template source using current registry without register it
    pub fn template_renderw2<T>(
        &self,
        template_source: &mut Read,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        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, data, writer)
    }
}

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

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

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

    static DUMMY_HELPER: DummyHelper = DummyHelper;

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

        assert!(r.register_template_string("index", "

").is_ok());
        assert!(r.register_template_string("index2", "

").is_ok());

        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(), 7 + 1);
    }

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

        assert!(r.register_template_string("index", "

").is_ok());

        let mut sw = StringWriter::new();
        {
            r.renderw("index", &(), &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_escape() {
        let r = Registry::new();
        let data = json!({
            "hello": "world"
        });

        assert_eq!(
            "{{hello}}",
            r.template_render(r"\\{{hello}}", &data).unwrap()
        );
    }
}
<commit_msg>(feat) allow source_map for template string render [#193]<commit_after>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;

use serde::Serialize;

use regex::{Captures, Regex};

use template::Template;
use render::{RenderContext, Renderable};
use context::Context;
use helpers::{self, HelperDef};
use directives::{self, DirectiveDef};
use support::str::StringWriter;
use error::{RenderError, 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.get(0).map(|m| m.as_str()) {
                Some("<") => "<",
                Some(">") => ">",
                Some("\"") => """,
                Some("&") => "&",
                _ => unreachable!(),
            }.to_owned()
        })
        .into_owned()
}

/// `EscapeFn` that do not 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()
    }

    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("log", Box::new(helpers::LOG_HELPER));

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

    /// Enable handlebars template source map
    ///
    /// Source map provides line/col reporting on error. It uses slightly
    /// more memory to maintain the data.
    ///
    /// Default is true.
    pub fn source_map_enable(&mut self, enable: bool) {
        self.source_map = enable;
    }

    /// Register a template string
    ///
    /// Returns `TemplateError` if there is syntax error on parsing template.
    pub fn register_template_string<S>(
        &mut self,
        name: &str,
        tpl_str: S,
    ) -> Result<(), TemplateError>
    where
        S: AsRef<str>,
    {
        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
    ///
    /// A named partial will be added to the registry. It will overwrite template with
    /// same name. Currently registered partial is just identical to template.
    pub fn register_partial<S>(&mut self, name: &str, partial_str: S) -> Result<(), TemplateError>
    where
        S: AsRef<str>,
    {
        self.register_template_string(name, partial_str)
    }

    /// Register a template from a path
    pub fn register_template_file
(
        &mut self,
        name: &str,
        tpl_path: P,
    ) -> Result<(), TemplateFileError>
    where
        P: AsRef<Path>,
    {
        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 decorator
    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
    ///
    /// * `name` is the template name you registred previously
    /// * `ctx` is the data that implements `serde::Serialize`
    ///
    /// Returns rendered string or an struct with error information
    pub fn render<T>(&self, name: &str, data: &T) -> Result<String, RenderError>
    where
        T: Serialize,
    {
        let mut writer = StringWriter::new();
        {
            try!(self.renderw(name, data, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// Render a registered template and write some data to the `std::io::Write`
    pub fn renderw<T>(&self, name: &str, data: &T, writer: &mut Write) -> Result<(), RenderError>
    where
        T: Serialize,
    {
        self.get_template(&name.to_string())
            .ok_or(RenderError::new(format!("Template not found: {}", name)))
            .and_then(|t| {
                let ctx = try!(Context::wraps(data));
                let mut local_helpers = HashMap::new();
                let mut render_context = RenderContext::new(ctx, &mut local_helpers, writer);
                render_context.root_template = t.name.clone();
                t.render(self, &mut render_context)
            })
    }

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

    /// render a template string using current registry without register it
    pub fn template_renderw<T>(
        &self,
        template_string: &str,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        let tpl = try!(Template::compile2(template_string, self.source_map));
        let ctx = try!(Context::wraps(data));
        let mut local_helpers = HashMap::new();
        let mut render_context = RenderContext::new(ctx, &mut local_helpers, writer);
        tpl.render(self, &mut render_context)
            .map_err(TemplateRenderError::from)
    }

    /// render a template source using current registry without register it
    pub fn template_renderw2<T>(
        &self,
        template_source: &mut Read,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        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, data, writer)
    }
}

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

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

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

    static DUMMY_HELPER: DummyHelper = DummyHelper;

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

        assert!(r.register_template_string("index", "

").is_ok());
        assert!(r.register_template_string("index2", "

").is_ok());

        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(), 7 + 1);
    }

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

        assert!(r.register_template_string("index", "

").is_ok());

        let mut sw = StringWriter::new();
        {
            r.renderw("index", &(), &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_escape() {
        let r = Registry::new();
        let data = json!({
            "hello": "world"
        });

        assert_eq!(
            "{{hello}}",
            r.template_render(r"\\{{hello}}", &data).unwrap()
        );
    }
}
<|endoftext|>


<commit_before>pub struct Entry {
    /// The pointer to the object stored in the registry. This is a type-erased
    /// `Box<T>`.
    pub ptr: *mut (),
    /// The function that can be called on the above pointer to drop the object
    /// and free its allocation.
    pub drop: unsafe fn(*mut ()),
}

#[cfg(feature = "slab")]
mod slab_impl {
    use std::cell::UnsafeCell;
    use std::num::NonZeroUsize;

    use super::Entry;

    struct Registry(slab::Slab<Entry>);

    impl Drop for Registry {
        fn drop(&mut self) {
            for (_, value) in self.0.iter() {
                // SAFETY: This function is only called once, and is called with the
                // pointer it was created with.
                unsafe { (value.drop)(value.ptr) };
            }
        }
    }

    thread_local!(static REGISTRY: UnsafeCell<Registry> = UnsafeCell::new(Registry(slab::Slab::new())));

    pub use usize as ItemId;

    pub fn insert(thread_id: NonZeroUsize, entry: Entry) -> ItemId {
        let _ = thread_id;

        // SAFETY: The `REGISTRY` is not accessed recursively in this function.
        REGISTRY.with(|registry| unsafe { (*registry.get()).0.insert(entry) })
    }

    pub fn with<R, F: FnOnce(&Entry) -> R>(item_id: ItemId, thread_id: NonZeroUsize, f: F) -> R {
        let _ = thread_id;
        REGISTRY.with(|registry| f(unsafe { &*registry.get() }.0.get(item_id).unwrap()))
    }

    pub fn remove(item_id: ItemId, thread_id: NonZeroUsize) -> Entry {
        let _ = thread_id;
        REGISTRY.with(|registry| unsafe { (*registry.get()).0.remove(item_id) })
    }
}

#[cfg(not(feature = "slab"))]
mod map_impl {
    use std::cell::UnsafeCell;
    use std::num::NonZeroUsize;

    use super::Entry;

    struct Registry(std::collections::HashMap<NonZeroUsize, Entry>);

    impl Drop for Registry {
        fn drop(&mut self) {
            for (_, value) in self.0.iter() {
                // SAFETY: This function is only called once, and is called with the
                // pointer it was created with.
                unsafe { (value.drop)(value.ptr) };
            }
        }
    }

    thread_local!(static REGISTRY: UnsafeCell<Registry> = UnsafeCell::new(Registry(Default::default())));

    pub type ItemId = ();

    pub fn insert(thread_id: NonZeroUsize, entry: Entry) -> ItemId {
        // SAFETY: The `REGISTRY` is not accessed recursively in this function.
        REGISTRY.with(|registry| unsafe { (*registry.get()).0.insert(thread_id, entry) });
    }

    pub fn with<R, F: FnOnce(&Entry) -> R>(item_id: ItemId, thread_id: NonZeroUsize, f: F) -> R {
        let _ = item_id;
        REGISTRY.with(|registry| f(unsafe { &*registry.get() }.0.get(&thread_id).unwrap()))
    }

    pub fn remove(item_id: ItemId, thread_id: NonZeroUsize) -> Entry {
        let _ = item_id;
        REGISTRY.with(|registry| unsafe { (*registry.get()).0.remove(&thread_id).unwrap() })
    }
}

#[cfg(feature = "slab")]
pub use self::slab_impl::*;

#[cfg(not(feature = "slab"))]
pub use self::map_impl::*;
<commit_msg>Small refactoring on the registry<commit_after>pub struct Entry {
    /// The pointer to the object stored in the registry. This is a type-erased
    /// `Box<T>`.
    pub ptr: *mut (),
    /// The function that can be called on the above pointer to drop the object
    /// and free its allocation.
    pub drop: unsafe fn(*mut ()),
}

#[cfg(feature = "slab")]
mod slab_impl {
    use std::cell::UnsafeCell;
    use std::num::NonZeroUsize;

    use super::Entry;

    pub struct Registry(pub slab::Slab<Entry>);

    thread_local!(static REGISTRY: UnsafeCell<Registry> = UnsafeCell::new(Registry(slab::Slab::new())));

    pub use usize as ItemId;

    pub fn insert(thread_id: NonZeroUsize, entry: Entry) -> ItemId {
        let _ = thread_id;
        REGISTRY.with(|registry| unsafe { (*registry.get()).0.insert(entry) })
    }

    pub fn with<R, F: FnOnce(&Entry) -> R>(item_id: ItemId, thread_id: NonZeroUsize, f: F) -> R {
        let _ = thread_id;
        REGISTRY.with(|registry| f(unsafe { &*registry.get() }.0.get(item_id).unwrap()))
    }

    pub fn remove(item_id: ItemId, thread_id: NonZeroUsize) -> Entry {
        let _ = thread_id;
        REGISTRY.with(|registry| unsafe { (*registry.get()).0.remove(item_id) })
    }
}

#[cfg(not(feature = "slab"))]
mod map_impl {
    use std::cell::UnsafeCell;
    use std::num::NonZeroUsize;

    use super::Entry;

    pub struct Registry(pub std::collections::HashMap<NonZeroUsize, Entry>);

    thread_local!(static REGISTRY: UnsafeCell<Registry> = UnsafeCell::new(Registry(Default::default())));

    pub type ItemId = ();

    pub fn insert(thread_id: NonZeroUsize, entry: Entry) -> ItemId {
        REGISTRY.with(|registry| unsafe { (*registry.get()).0.insert(thread_id, entry) });
    }

    pub fn with<R, F: FnOnce(&Entry) -> R>(item_id: ItemId, thread_id: NonZeroUsize, f: F) -> R {
        let _ = item_id;
        REGISTRY.with(|registry| f(unsafe { &*registry.get() }.0.get(&thread_id).unwrap()))
    }

    pub fn remove(item_id: ItemId, thread_id: NonZeroUsize) -> Entry {
        let _ = item_id;
        REGISTRY.with(|registry| unsafe { (*registry.get()).0.remove(&thread_id).unwrap() })
    }
}

#[cfg(feature = "slab")]
pub use self::slab_impl::*;

#[cfg(not(feature = "slab"))]
pub use self::map_impl::*;

impl Drop for Registry {
    fn drop(&mut self) {
        for (_, value) in self.0.iter() {
            // SAFETY: This function is only called once, and is called with the
            // pointer it was created with.
            unsafe { (value.drop)(value.ptr) };
        }
    }
}
<|endoftext|>


<commit_before>// Type for a syntax tree node that is reserved for future use.
//
// For example ExprReference contains a field `raw` of type Reserved. If `&raw
// place` syntax becomes a thing as per https://github.com/rust-lang/rfcs/pull/2582,
// we can backward compatibly change `raw`'s type to Option<Token![raw]> without
// the possibility of breaking any code.

ast_struct! {
    #[derive(Default)]
    pub struct Reserved {
        _private: (),
    }
}
<commit_msg>Ensure appropriate autotrait impls for Reserved as if it contained Span<commit_after>// Type for a syntax tree node that is reserved for future use.
//
// For example ExprReference contains a field `raw` of type Reserved. If `&raw
// place` syntax becomes a thing as per https://github.com/rust-lang/rfcs/pull/2582,
// we can backward compatibly change `raw`'s type to Option<Token![raw]> without
// the possibility of breaking any code.

use proc_macro2::Span;
use std::marker::PhantomData;

ast_struct! {
    #[derive(Default)]
    pub struct Reserved {
        _private: PhantomData<Span>,
    }
}
<|endoftext|>


<commit_before>/* CONFIGURATION */

// add other android system headers to this list as necessary
static INCLUDE_SYS_H: [&str; 1] = ["jni.h"];

static RUST_SRC_DIR: &str = "src";
static ANDROID_BASE_DIR: &str = "app";

// =========================================================
// This script is portable; copy and paste it into your own
// build files at will.

extern crate bindgen;
extern crate rust_swig;
extern crate syntex;
extern crate walkdir;

use std::process::{Command, Stdio};
use std::{env, fmt};
use std::io::prelude::*;
use std::fs::File;
use std::path::{Path, PathBuf};

use bindgen::RustTarget;
use rust_swig::{JavaConfig, LanguageConfig};
use walkdir::WalkDir;

fn main() {
    // don't simplify this to if the target contains the substring "android" --
    // these lines also serve as a guard so only true android triples receive
    // JNI generation.
    let target = env::var("TARGET").unwrap();
    if ["aarch64-linux-android",
        "arm-linux-androideabi",
        "i686-linux-android",
        "x86_64-linux-android"]
        .contains(&target.as_str()) {
            gen_for_android();
        }
}

fn gen_for_android() {
    let target = env::var("TARGET").unwrap();

    let include_dirs =
        get_gcc_system_include_dirs(&target).expect("Can't get NDK's system include dirs");

    let include_headers: Vec<_> = INCLUDE_SYS_H.into_iter()
        .map(|h| search_file_in_directory(&include_dirs, h)
             .expect(format!("Could not find header {}", h).as_ref()))
        .collect();

    let src_dir = Path::new(RUST_SRC_DIR);
    let out_dir = env::var("OUT_DIR").unwrap();

    gen_binding(
        &target,
        &include_dirs,
        &include_headers,
        &Path::new(&out_dir).join("android_c_headers.rs"),
    ).unwrap();

    // Find files ending in .rs.in and expand them with SWIG
    for _entry in WalkDir::new(RUST_SRC_DIR) {
        let entry = _entry.expect("Error walking sources.");
        if entry.path().is_dir() || !entry.path().to_string_lossy().ends_with(".rs.in") {
            continue;
        }

        println!("Found SWIG specification: {}", entry.path().display());
        let swigf = entry.path().strip_prefix("src").unwrap();

        rust_swig_expand(
            &src_dir,
            &swigf,
            Path::new(&out_dir),
        ).unwrap();

        write_include_file(&src_dir, swigf).expect("Failed to write include file.");
    }

    // Hook up cargo reruns
    println!("cargo:rerun-if-changed={}", RUST_SRC_DIR);
    for dir in &include_dirs {println!("cargo:rerun-if-changed={}", dir.display());}
    //if the generated files were deleted (e.g by gradle -q clean), regenerate them
    println!("cargo:rerun-if-changed={}", out_dir);
}

fn get_gcc_system_include_dirs(target: &str) -> Result<Vec<PathBuf>, String> {
    let gcc_cmd = match env::var("RUSTC_LINKER") {
        Ok(path) => path,
        Err(_) => target.to_owned() + "-gcc"
    };

    assert!(Path::new(&gcc_cmd).exists(),
            "RUSTC_LINKER not set -- Update your version of Cargo
or merge PR#5394 from github.com/rust-lang/cargo.");

    println!("Using Android gcc from '{}'", gcc_cmd);

    let gcc_process = Command::new(&gcc_cmd)
        .args(&["-v", "-x", "c", "-E", "-" /*, &("--sysroot=".to_owned() + &sysroot)*/])
        .stderr(Stdio::piped())
        .stdin(Stdio::piped())
        .stdout(Stdio::inherit())
        .spawn()
        .map_err(|err| err.to_string())?;

    gcc_process
        .stdin
        .ok_or(format!("Cannot get stdin of {}", gcc_cmd).as_str())?
        .write_all("\n".as_bytes())
        .map_err(|err| err.to_string())?;

    let mut gcc_output = String::new();
    gcc_process
        .stderr
        .ok_or(format!("Cannot get stderr of {}", gcc_cmd).as_str())?
        .read_to_string(&mut gcc_output)
        .map_err(|err| err.to_string())?;

    const BEGIN_PAT: &'static str = "\n#include <...> search starts here:\n";
    const END_PAT: &'static str = "\nEnd of search list.\n";
    let start_includes = gcc_output
        .find(BEGIN_PAT)
        .ok_or(format!("No '{}' in output from {}", BEGIN_PAT, gcc_cmd).as_str())?
        + BEGIN_PAT.len();
    let end_includes = (&gcc_output[start_includes..])
        .find(END_PAT)
        .ok_or(format!("No '{}' in output from {}", END_PAT, gcc_cmd).as_str())?
        + start_includes;

    Ok((&gcc_output[start_includes..end_includes])
       .split('\n')
       .map(|s| PathBuf::from(s.trim().to_string()))
       .collect())
}

fn search_file_in_directory<P>(dirs: &[P], file: &str) -> Result<PathBuf, ()>
where
    P: AsRef<Path>,
{
    for dir in dirs {
        let file_path = dir.as_ref().join(file);
        if file_path.exists() && file_path.is_file() {
            return Ok(file_path);
        }
    }
    Err(())
}

fn gen_binding<P1, P2>(
    target: &str,
    include_dirs: &[P1],
    c_headers: &[P2],
    output_rust: &Path,
) -> Result<(), String>
where
    P1: AsRef<Path> + fmt::Debug,
    P2: AsRef<Path> + fmt::Debug,
{
    assert!(!c_headers.is_empty());
    let c_file_path = &c_headers[0];

    let mut bindings: bindgen::Builder =
        bindgen::builder().header(c_file_path.as_ref().to_str().unwrap());
    bindings = include_dirs.iter().fold(bindings, |acc, x| {
        acc.clang_arg("-I".to_string() + x.as_ref().to_str().unwrap())
    });
    println!("Generate binding for {:?}", c_headers);
    bindings = bindings
        .rust_target(RustTarget::Stable_1_19)
    //long double not supported yet, see https://github.com/servo/rust-bindgen/issues/550
        .blacklist_type("max_align_t");
    bindings = if target.contains("windows") {
        //see https://github.com/servo/rust-bindgen/issues/578
        bindings.trust_clang_mangling(false)
    } else {
        bindings
    };
    bindings = c_headers[1..].iter().fold(
        Ok(bindings),
        |acc: Result<bindgen::Builder, String>, header| {
            let c_file_path = header;
            let c_file_str = c_file_path
                .as_ref()
                .to_str()
                .ok_or_else(|| format!("Invalid unicode in path to {:?}", c_file_path.as_ref()))?;
            Ok(acc.unwrap().clang_arg("-include").clang_arg(c_file_str))
        },
    )?;

    let generated_bindings = bindings
    //        .clang_arg(format!("-target {}", target))
        .generate()
        .map_err(|_| "Failed to generate bindings".to_string())?;
    generated_bindings
        .write_to_file(output_rust)
        .map_err(|err| err.to_string())?;

    Ok(())
}

fn rust_swig_expand(source_dir: &Path, file: &Path, out_dir: &Path) -> Result<(), String> {
    let app_id = env::var("ANDROID_APPLICATION_ID")
        .expect("You must set ANDROID_APPLICATION_ID to the package name of your android application.");

    let mut registry = syntex::Registry::new();
    let swig_gen = rust_swig::Generator::new(LanguageConfig::JavaConfig(
        JavaConfig::new(
            Path::new(ANDROID_BASE_DIR)
                .join("src")
                .join("main")
                .join("java")
                .join(app_id.replace(".", "/")),
            app_id.to_owned() // not sure why it requires ownership...
        ).use_null_annotation("android.support.annotation.NonNull".into()),
    ));
    swig_gen.register(&mut registry);

    let out_file = out_dir.join(Path::new(file.parent().unwrap_or(Path::new(".")))
                                .join(file.file_stem().expect("Got invalid file (no filename)")));

    registry
        .expand(app_id.as_ref(), source_dir.join(file), out_file)
        .map_err(|err| format!("Rust-SWIG expansion failed: {}", err))
}

fn write_include_file(source_dir: &Path, swig_file: &Path) -> std::io::Result<()> {
    let rs_rel_file = Path::new(swig_file.parent().unwrap_or(Path::new(".")))
        .join(swig_file.file_stem().expect("Got invalid file (no filename)"));
    let rs_path = source_dir.join(&rs_rel_file);

    if rs_path.exists() {
        println!("Not writing {} because it exists", rs_path.display());
        return Ok(());
    }

    let mut rs_file = File::create(rs_path)?;
    rs_file.write_all(format!("include!(concat!(env!(\"OUT_DIR\"), \"/{}\"",
                              rs_rel_file.to_string_lossy()).as_bytes())?;

    return Ok(());
}
<commit_msg>Improved error message on unfound linker<commit_after>/* CONFIGURATION */

// add other android system headers to this list as necessary
static INCLUDE_SYS_H: [&str; 1] = ["jni.h"];

static RUST_SRC_DIR: &str = "src";
static ANDROID_BASE_DIR: &str = "app";

// =========================================================
// This script is portable; copy and paste it into your own
// build files at will.

extern crate bindgen;
extern crate rust_swig;
extern crate syntex;
extern crate walkdir;

use std::process::{Command, Stdio};
use std::{env, fmt};
use std::io::prelude::*;
use std::fs::File;
use std::path::{Path, PathBuf};

use bindgen::RustTarget;
use rust_swig::{JavaConfig, LanguageConfig};
use walkdir::WalkDir;

fn main() {
    // don't simplify this to if the target contains the substring "android" --
    // these lines also serve as a guard so only true android triples receive
    // JNI generation.
    let target = env::var("TARGET").unwrap();
    if ["aarch64-linux-android",
        "arm-linux-androideabi",
        "i686-linux-android",
        "x86_64-linux-android"]
        .contains(&target.as_str()) {
            gen_for_android();
        }
}

fn gen_for_android() {
    let target = env::var("TARGET").unwrap();

    let include_dirs =
        get_gcc_system_include_dirs(&target).expect("Can't get NDK's system include dirs");

    let include_headers: Vec<_> = INCLUDE_SYS_H.into_iter()
        .map(|h| search_file_in_directory(&include_dirs, h)
             .expect(format!("Could not find header {}", h).as_ref()))
        .collect();

    let src_dir = Path::new(RUST_SRC_DIR);
    let out_dir = env::var("OUT_DIR").unwrap();

    gen_binding(
        &target,
        &include_dirs,
        &include_headers,
        &Path::new(&out_dir).join("android_c_headers.rs"),
    ).unwrap();

    // Find files ending in .rs.in and expand them with SWIG
    for _entry in WalkDir::new(RUST_SRC_DIR) {
        let entry = _entry.expect("Error walking sources.");
        if entry.path().is_dir() || !entry.path().to_string_lossy().ends_with(".rs.in") {
            continue;
        }

        println!("Found SWIG specification: {}", entry.path().display());
        let swigf = entry.path().strip_prefix("src").unwrap();

        rust_swig_expand(
            &src_dir,
            &swigf,
            Path::new(&out_dir),
        ).unwrap();

        write_include_file(&src_dir, swigf).expect("Failed to write include file.");
    }

    // Hook up cargo reruns
    println!("cargo:rerun-if-changed={}", RUST_SRC_DIR);
    for dir in &include_dirs {println!("cargo:rerun-if-changed={}", dir.display());}
    //if the generated files were deleted (e.g by gradle -q clean), regenerate them
    println!("cargo:rerun-if-changed={}", out_dir);
}

fn get_gcc_system_include_dirs(target: &str) -> Result<Vec<PathBuf>, String> {
    let gcc_cmd = match env::var("RUSTC_LINKER") {
        Ok(path) => path,
        Err(_) => target.to_owned() + "-gcc"
    };

    let gcc_path = Path::new(&gcc_cmd);
    assert!(gcc_path.exists() && gcc_path.is_file(), format!(
            "Could not find a suitable NDK gcc (tried {})
    You can fix this either by adding the toolchain's bin/ to your $PATH, or by
    merging PR#5394 form github.com/rust-lang/cargo into your version of Cargo.", &gcc_cmd));
    // TODO: Remove merge instruction once PR is merged into stable cargo.

    println!("Using Android gcc from '{}'", gcc_cmd);

    let gcc_process = Command::new(&gcc_cmd)
        .args(&["-v", "-x", "c", "-E", "-"])
        .stderr(Stdio::piped())
        .stdin(Stdio::piped())
        .stdout(Stdio::inherit())
        .spawn()
        .map_err(|err| err.to_string())?;

    gcc_process
        .stdin
        .ok_or(format!("Cannot get stdin of {}", gcc_cmd).as_str())?
        .write_all("\n".as_bytes())
        .map_err(|err| err.to_string())?;

    let mut gcc_output = String::new();
    gcc_process
        .stderr
        .ok_or(format!("Cannot get stderr of {}", gcc_cmd).as_str())?
        .read_to_string(&mut gcc_output)
        .map_err(|err| err.to_string())?;

    const BEGIN_PAT: &'static str = "\n#include <...> search starts here:\n";
    const END_PAT: &'static str = "\nEnd of search list.\n";
    let start_includes = gcc_output
        .find(BEGIN_PAT)
        .ok_or(format!("No '{}' in output from {}", BEGIN_PAT, gcc_cmd).as_str())?
        + BEGIN_PAT.len();
    let end_includes = (&gcc_output[start_includes..])
        .find(END_PAT)
        .ok_or(format!("No '{}' in output from {}", END_PAT, gcc_cmd).as_str())?
        + start_includes;

    Ok((&gcc_output[start_includes..end_includes])
       .split('\n')
       .map(|s| PathBuf::from(s.trim().to_string()))
       .collect())
}

fn search_file_in_directory<P>(dirs: &[P], file: &str) -> Result<PathBuf, ()>
where
    P: AsRef<Path>,
{
    for dir in dirs {
        let file_path = dir.as_ref().join(file);
        if file_path.exists() && file_path.is_file() {
            return Ok(file_path);
        }
    }
    Err(())
}

fn gen_binding<P1, P2>(
    target: &str,
    include_dirs: &[P1],
    c_headers: &[P2],
    output_rust: &Path,
) -> Result<(), String>
where
    P1: AsRef<Path> + fmt::Debug,
    P2: AsRef<Path> + fmt::Debug,
{
    assert!(!c_headers.is_empty());
    let c_file_path = &c_headers[0];

    let mut bindings: bindgen::Builder =
        bindgen::builder().header(c_file_path.as_ref().to_str().unwrap());
    bindings = include_dirs.iter().fold(bindings, |acc, x| {
        acc.clang_arg("-I".to_string() + x.as_ref().to_str().unwrap())
    });
    println!("Generate binding for {:?}", c_headers);
    bindings = bindings
        .rust_target(RustTarget::Stable_1_19)
    //long double not supported yet, see https://github.com/servo/rust-bindgen/issues/550
        .blacklist_type("max_align_t");
    bindings = if target.contains("windows") {
        //see https://github.com/servo/rust-bindgen/issues/578
        bindings.trust_clang_mangling(false)
    } else {
        bindings
    };
    bindings = c_headers[1..].iter().fold(
        Ok(bindings),
        |acc: Result<bindgen::Builder, String>, header| {
            let c_file_path = header;
            let c_file_str = c_file_path
                .as_ref()
                .to_str()
                .ok_or_else(|| format!("Invalid unicode in path to {:?}", c_file_path.as_ref()))?;
            Ok(acc.unwrap().clang_arg("-include").clang_arg(c_file_str))
        },
    )?;

    let generated_bindings = bindings
    //        .clang_arg(format!("-target {}", target))
        .generate()
        .map_err(|_| "Failed to generate bindings".to_string())?;
    generated_bindings
        .write_to_file(output_rust)
        .map_err(|err| err.to_string())?;

    Ok(())
}

fn rust_swig_expand(source_dir: &Path, file: &Path, out_dir: &Path) -> Result<(), String> {
    let app_id = env::var("ANDROID_APPLICATION_ID")
        .expect("You must set ANDROID_APPLICATION_ID to the package name of your android application.");

    let mut registry = syntex::Registry::new();
    let swig_gen = rust_swig::Generator::new(LanguageConfig::JavaConfig(
        JavaConfig::new(
            Path::new(ANDROID_BASE_DIR)
                .join("src")
                .join("main")
                .join("java")
                .join(app_id.replace(".", "/")),
            app_id.to_owned() // not sure why it requires ownership...
        ).use_null_annotation("android.support.annotation.NonNull".into()),
    ));
    swig_gen.register(&mut registry);

    let out_file = out_dir.join(Path::new(file.parent().unwrap_or(Path::new(".")))
                                .join(file.file_stem().expect("Got invalid file (no filename)")));

    registry
        .expand(app_id.as_ref(), source_dir.join(file), out_file)
        .map_err(|err| format!("Rust-SWIG expansion failed: {}", err))
}

fn write_include_file(source_dir: &Path, swig_file: &Path) -> std::io::Result<()> {
    let rs_rel_file = Path::new(swig_file.parent().unwrap_or(Path::new(".")))
        .join(swig_file.file_stem().expect("Got invalid file (no filename)"));
    let rs_path = source_dir.join(&rs_rel_file);

    if rs_path.exists() {
        println!("Not writing {} because it exists", rs_path.display());
        return Ok(());
    }

    let mut rs_file = File::create(rs_path)?;
    rs_file.write_all(format!("include!(concat!(env!(\"OUT_DIR\"), \"/{}\"",
                              rs_rel_file.to_string_lossy()).as_bytes())?;

    return Ok(());
}
<|endoftext|>


<commit_before>//! Access Control List (ACL) for shadowsocks
//!
//! This is for advance controlling server behaviors in both local and proxy servers.

use std::{
    fmt,
    fs::File,
    io::{self, BufRead, BufReader, Error, ErrorKind},
    net::{IpAddr, SocketAddr},
    path::Path,
};

use ipnet::{IpNet, Ipv4Net, Ipv6Net};
use iprange::IpRange;
use regex::{RegexSet, RegexSetBuilder};
use trust_dns_proto::{
    op::Query,
    rr::RecordType,
};

use crate::{context::Context, relay::socks5::Address};

/// Strategy mode that ACL is running
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Mode {
    /// BlackList mode, rejects or bypasses all requests by default
    BlackList,
    /// WhiteList mode, accepts or proxies all requests by default
    WhiteList,
}

#[derive(Clone)]
struct Rules {
    ipv4: IpRange<Ipv4Net>,
    ipv6: IpRange<Ipv6Net>,
    rule: RegexSet,
}

impl fmt::Debug for Rules {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Rules {{ ipv4: {:?}, ipv6: {:?}, rule: [", self.ipv4, self.ipv6)?;

        let max_len = 2;
        let has_more = self.rule.len() > max_len;

        for (idx, r) in self.rule.patterns().iter().take(max_len).enumerate() {
            if idx > 0 {
                f.write_str(", ")?;
            }
            f.write_str(r)?;
        }

        if has_more {
            f.write_str(", ...")?;
        }

        f.write_str("] }")
    }
}

impl Rules {
    /// Create a new rule
    fn new(mut ipv4: IpRange<Ipv4Net>, mut ipv6: IpRange<Ipv6Net>, rule: RegexSet) -> Rules {
        // Optimization, merging networks
        ipv4.simplify();
        ipv6.simplify();

        Rules { ipv4, ipv6, rule }
    }

    /// Check if the specified address matches these rules
    fn check_address_matched(&self, addr: &Address) -> bool {
        match *addr {
            Address::SocketAddress(ref saddr) => self.check_ip_matched(&saddr.ip()),
            Address::DomainNameAddress(ref domain, ..) => self.check_host_matched(domain),
        }
    }

    /// Check if the specified address matches any rules
    fn check_ip_matched(&self, addr: &IpAddr) -> bool {
        match addr {
            IpAddr::V4(v4) => self.ipv4.contains(v4),
            IpAddr::V6(v6) => self.ipv6.contains(v6),
        }
    }

    /// Check if the specified host matches any rules
    fn check_host_matched(&self, host: &str) -> bool {
        self.rule.is_match(host)
    }

    /// Check if there are no rules for the corresponding DNS query type
    fn is_rule_empty_for_qtype(&self, qtype: RecordType) -> bool {
        match qtype {
            RecordType::A => self.ipv4.iter().next().is_none(),
            RecordType::AAAA => self.ipv6.iter().next().is_none(),
            _ => true
        }
    }
}

/// ACL rules
///
/// ## Sections
///
/// ACL File is formatted in sections, each section has a name with surrounded by brackets `[` and `]`
/// followed by Rules line by line.
///
/// ```plain
/// [SECTION-1]
/// RULE-1
/// RULE-2
/// RULE-3
///
/// [SECTION-2]
/// RULE-1
/// RULE-2
/// RULE-3
/// ```
///
/// Available sections are
///
/// - For local servers (`sslocal`, `ssredir`, ...)
///     * `[bypass_all]` - ACL runs in `BlackList` mode.
///     * `[proxy_all]` - ACL runs in `WhiteList` mode.
///     * `[bypass_list]` - Rules for connecting directly
///     * `[proxy_list]` - Rules for connecting through proxies
/// - For remote servers (`ssserver`)
///     * `[reject_all]` - ACL runs in `BlackList` mode.
///     * `[accept_all]` - ACL runs in `WhiteList` mode.
///     * `[black_list]` - Rules for rejecting
///     * `[white_list]` - Rules for allowing
///     * `[outbound_block_list]` - Rules for blocking outbound addresses.
///
/// ## Mode
///
/// Mode is the default ACL strategy for those addresses that are not in configuration file.
///
/// - `BlackList` - Bypasses / Rejects all addresses except those in `[proxy_list]` or `[white_list]`
/// - `WhiltList` - Proxies / Accepts all addresses except those in `[bypass_list]` or `[black_list]`
///
/// ## Rules
///
/// Rules can be either
///
/// - CIDR form network addresses, like `10.9.0.32/16`
/// - IP addresses, like `127.0.0.1` or `::1`
/// - Regular Expression for matching hosts, like `(^|\.)gmail\.com$`
#[derive(Debug, Clone)]
pub struct AccessControl {
    outbound_block: Rules,
    black_list: Rules,
    white_list: Rules,
    mode: Mode,
}

impl AccessControl {
    /// Load ACL rules from a file
    pub fn load_from_file<P: AsRef<Path>>(p: P) -> io::Result<AccessControl> {
        let fp = File::open(p)?;
        let r = BufReader::new(fp);

        let mut mode = Mode::BlackList;

        let mut outbound_block_ipv4 = IpRange::new();
        let mut outbound_block_ipv6 = IpRange::new();
        let mut outbound_block_rules = Vec::new();
        let mut bypass_ipv4 = IpRange::new();
        let mut bypass_ipv6 = IpRange::new();
        let mut bypass_rules = Vec::new();
        let mut proxy_ipv4 = IpRange::new();
        let mut proxy_ipv6 = IpRange::new();
        let mut proxy_rules = Vec::new();

        let mut curr_ipv4 = &mut bypass_ipv4;
        let mut curr_ipv6 = &mut bypass_ipv6;
        let mut curr_rules = &mut bypass_rules;

        for line in r.lines() {
            let line = line?;
            if line.is_empty() {
                continue;
            }

            // Comments
            if line.starts_with('#') {
                continue;
            }

            match line.as_str() {
                "[reject_all]" | "[bypass_all]" => {
                    mode = Mode::WhiteList;
                }
                "[accept_all]" | "[proxy_all]" => {
                    mode = Mode::BlackList;
                }
                "[outbound_block_list]" => {
                    curr_ipv4 = &mut outbound_block_ipv4;
                    curr_ipv6 = &mut outbound_block_ipv6;
                    curr_rules = &mut outbound_block_rules;
                }
                "[black_list]" | "[bypass_list]" => {
                    curr_ipv4 = &mut bypass_ipv4;
                    curr_ipv6 = &mut bypass_ipv6;
                    curr_rules = &mut bypass_rules;
                }
                "[white_list]" | "[proxy_list]" => {
                    curr_ipv4 = &mut proxy_ipv4;
                    curr_ipv6 = &mut proxy_ipv6;
                    curr_rules = &mut proxy_rules;
                }
                _ => {
                    match line.parse::<IpNet>() {
                        Ok(IpNet::V4(v4)) => {
                            curr_ipv4.add(v4);
                        }
                        Ok(IpNet::V6(v6)) => {
                            curr_ipv6.add(v6);
                        }
                        Err(..) => {
                            // Maybe it is a pure IpAddr
                            match line.parse::<IpAddr>() {
                                Ok(IpAddr::V4(v4)) => {
                                    curr_ipv4.add(Ipv4Net::from(v4));
                                }
                                Ok(IpAddr::V6(v6)) => {
                                    curr_ipv6.add(Ipv6Net::from(v6));
                                }
                                Err(..) => {
                                    // FIXME: If this line is not a valid regex, how can we know without actually compile it?
                                    curr_rules.push(line);
                                }
                            }
                        }
                    }
                }
            }
        }

        const REGEX_SIZE_LIMIT: usize = usize::max_value();

        let outbound_block_regex = match RegexSetBuilder::new(outbound_block_rules)
            .size_limit(REGEX_SIZE_LIMIT)
            .build()
        {
            Ok(r) => r,
            Err(err) => {
                let err = Error::new(ErrorKind::Other, format!("[outbound_block_list] regex error: {}", err));
                return Err(err);
            }
        };

        let bypass_regex = match RegexSetBuilder::new(bypass_rules).size_limit(REGEX_SIZE_LIMIT).build() {
            Ok(r) => r,
            Err(err) => {
                let err = Error::new(
                    ErrorKind::Other,
                    format!("[black_list] or [bypass_list] regex error: {}", err),
                );
                return Err(err);
            }
        };

        let proxy_regex = match RegexSetBuilder::new(proxy_rules).size_limit(REGEX_SIZE_LIMIT).build() {
            Ok(r) => r,
            Err(err) => {
                let err = Error::new(
                    ErrorKind::Other,
                    format!("[white_list] or [proxy_list] regex error: {}", err),
                );
                return Err(err);
            }
        };

        Ok(AccessControl {
            outbound_block: Rules::new(outbound_block_ipv4, outbound_block_ipv6, outbound_block_regex),
            black_list: Rules::new(bypass_ipv4, bypass_ipv6, bypass_regex),
            white_list: Rules::new(proxy_ipv4, proxy_ipv6, proxy_regex),
            mode,
        })
    }

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    pub fn check_query_in_proxy_list(&self, query: &Query) -> Option<bool> {
        // remove the last dot from fqdn name
        let mut name = query.name().to_ascii();
        name.pop();
        let addr = Address::DomainNameAddress(name, 0);
        // Addresses in proxy_list will be proxied
        if self.white_list.check_address_matched(&addr) {
            return Some(true);
        }
        if self.black_list.check_address_matched(&addr) {
            return Some(false);
        }
        match self.mode {
            Mode::BlackList => if self.black_list.is_rule_empty_for_qtype(query.query_type()) {
                return Some(true);
            },
            Mode::WhiteList => if self.white_list.is_rule_empty_for_qtype(query.query_type()) {
                return Some(false);
            },
        }
        None
    }

    pub fn check_ip_in_proxy_list(&self, ip: &IpAddr) -> bool {
        match self.mode {
            Mode::BlackList => !self.black_list.check_ip_matched(ip),
            Mode::WhiteList => self.white_list.check_ip_matched(ip),
        }
    }

    /// Check if target address should be bypassed (for client)
    ///
    /// FIXME: This function may perform a DNS resolution
    pub async fn check_target_bypassed(&self, context: &Context, addr: &Address) -> bool {
        // Addresses in bypass_list will be bypassed
        if self.black_list.check_address_matched(addr) {
            return true;
        }

        // Addresses in proxy_list will be proxied
        if self.white_list.check_address_matched(addr) {
            return false;
        }

        // Resolve hostname and check the list
        if cfg!(not(target_os = "android")) {
            if let Address::DomainNameAddress(ref host, port) = *addr {
                if let Ok(vaddr) = context.dns_resolve(host, port).await {
                    for addr in vaddr {
                        if self.black_list.check_ip_matched(&addr.ip()) {
                            return true;
                        }

                        if self.white_list.check_ip_matched(&addr.ip()) {
                            return false;
                        }
                    }
                }
            }
        }

        // default rule
        match self.mode {
            Mode::BlackList => false,
            Mode::WhiteList => true,
        }
    }

    /// Check if client address should be blocked (for server)
    pub fn check_client_blocked(&self, addr: &SocketAddr) -> bool {
        match self.mode {
            Mode::BlackList => {
                // Only clients in black_list will be blocked
                self.black_list.check_ip_matched(&addr.ip())
            }
            Mode::WhiteList => {
                // Only clients in white_list will be proxied
                !self.white_list.check_ip_matched(&addr.ip())
            }
        }
    }

    /// Check if outbound address is blocked (for server)
    ///
    /// NOTE: `Address::DomainName` is only validated by regex rules,
    ///       resolved addresses are checked in the `lookup_outbound_then!` macro
    pub fn check_outbound_blocked(&self, outbound: &Address) -> bool {
        self.outbound_block.check_address_matched(outbound)
    }

    /// Check resolved outbound address is blocked (for server)
    pub fn check_resolved_outbound_blocked(&self, outbound: &SocketAddr) -> bool {
        self.outbound_block.check_ip_matched(&outbound.ip())
    }
}
<commit_msg>Filter out non-IN queries<commit_after>//! Access Control List (ACL) for shadowsocks
//!
//! This is for advance controlling server behaviors in both local and proxy servers.

use std::{
    fmt,
    fs::File,
    io::{self, BufRead, BufReader, Error, ErrorKind},
    net::{IpAddr, SocketAddr},
    path::Path,
};

use ipnet::{IpNet, Ipv4Net, Ipv6Net};
use iprange::IpRange;
use regex::{RegexSet, RegexSetBuilder};
use trust_dns_proto::{
    op::Query,
    rr::{DNSClass, RecordType},
};

use crate::{context::Context, relay::socks5::Address};

/// Strategy mode that ACL is running
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Mode {
    /// BlackList mode, rejects or bypasses all requests by default
    BlackList,
    /// WhiteList mode, accepts or proxies all requests by default
    WhiteList,
}

#[derive(Clone)]
struct Rules {
    ipv4: IpRange<Ipv4Net>,
    ipv6: IpRange<Ipv6Net>,
    rule: RegexSet,
}

impl fmt::Debug for Rules {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Rules {{ ipv4: {:?}, ipv6: {:?}, rule: [", self.ipv4, self.ipv6)?;

        let max_len = 2;
        let has_more = self.rule.len() > max_len;

        for (idx, r) in self.rule.patterns().iter().take(max_len).enumerate() {
            if idx > 0 {
                f.write_str(", ")?;
            }
            f.write_str(r)?;
        }

        if has_more {
            f.write_str(", ...")?;
        }

        f.write_str("] }")
    }
}

impl Rules {
    /// Create a new rule
    fn new(mut ipv4: IpRange<Ipv4Net>, mut ipv6: IpRange<Ipv6Net>, rule: RegexSet) -> Rules {
        // Optimization, merging networks
        ipv4.simplify();
        ipv6.simplify();

        Rules { ipv4, ipv6, rule }
    }

    /// Check if the specified address matches these rules
    fn check_address_matched(&self, addr: &Address) -> bool {
        match *addr {
            Address::SocketAddress(ref saddr) => self.check_ip_matched(&saddr.ip()),
            Address::DomainNameAddress(ref domain, ..) => self.check_host_matched(domain),
        }
    }

    /// Check if the specified address matches any rules
    fn check_ip_matched(&self, addr: &IpAddr) -> bool {
        match addr {
            IpAddr::V4(v4) => self.ipv4.contains(v4),
            IpAddr::V6(v6) => self.ipv6.contains(v6),
        }
    }

    /// Check if the specified host matches any rules
    fn check_host_matched(&self, host: &str) -> bool {
        self.rule.is_match(host)
    }

    /// Check if there are no rules for the corresponding DNS query type
    fn is_rule_empty_for_qtype(&self, qtype: RecordType) -> bool {
        match qtype {
            RecordType::A => self.ipv4.iter().next().is_none(),
            RecordType::AAAA => self.ipv6.iter().next().is_none(),
            _ => true
        }
    }
}

/// ACL rules
///
/// ## Sections
///
/// ACL File is formatted in sections, each section has a name with surrounded by brackets `[` and `]`
/// followed by Rules line by line.
///
/// ```plain
/// [SECTION-1]
/// RULE-1
/// RULE-2
/// RULE-3
///
/// [SECTION-2]
/// RULE-1
/// RULE-2
/// RULE-3
/// ```
///
/// Available sections are
///
/// - For local servers (`sslocal`, `ssredir`, ...)
///     * `[bypass_all]` - ACL runs in `BlackList` mode.
///     * `[proxy_all]` - ACL runs in `WhiteList` mode.
///     * `[bypass_list]` - Rules for connecting directly
///     * `[proxy_list]` - Rules for connecting through proxies
/// - For remote servers (`ssserver`)
///     * `[reject_all]` - ACL runs in `BlackList` mode.
///     * `[accept_all]` - ACL runs in `WhiteList` mode.
///     * `[black_list]` - Rules for rejecting
///     * `[white_list]` - Rules for allowing
///     * `[outbound_block_list]` - Rules for blocking outbound addresses.
///
/// ## Mode
///
/// Mode is the default ACL strategy for those addresses that are not in configuration file.
///
/// - `BlackList` - Bypasses / Rejects all addresses except those in `[proxy_list]` or `[white_list]`
/// - `WhiltList` - Proxies / Accepts all addresses except those in `[bypass_list]` or `[black_list]`
///
/// ## Rules
///
/// Rules can be either
///
/// - CIDR form network addresses, like `10.9.0.32/16`
/// - IP addresses, like `127.0.0.1` or `::1`
/// - Regular Expression for matching hosts, like `(^|\.)gmail\.com$`
#[derive(Debug, Clone)]
pub struct AccessControl {
    outbound_block: Rules,
    black_list: Rules,
    white_list: Rules,
    mode: Mode,
}

impl AccessControl {
    /// Load ACL rules from a file
    pub fn load_from_file<P: AsRef<Path>>(p: P) -> io::Result<AccessControl> {
        let fp = File::open(p)?;
        let r = BufReader::new(fp);

        let mut mode = Mode::BlackList;

        let mut outbound_block_ipv4 = IpRange::new();
        let mut outbound_block_ipv6 = IpRange::new();
        let mut outbound_block_rules = Vec::new();
        let mut bypass_ipv4 = IpRange::new();
        let mut bypass_ipv6 = IpRange::new();
        let mut bypass_rules = Vec::new();
        let mut proxy_ipv4 = IpRange::new();
        let mut proxy_ipv6 = IpRange::new();
        let mut proxy_rules = Vec::new();

        let mut curr_ipv4 = &mut bypass_ipv4;
        let mut curr_ipv6 = &mut bypass_ipv6;
        let mut curr_rules = &mut bypass_rules;

        for line in r.lines() {
            let line = line?;
            if line.is_empty() {
                continue;
            }

            // Comments
            if line.starts_with('#') {
                continue;
            }

            match line.as_str() {
                "[reject_all]" | "[bypass_all]" => {
                    mode = Mode::WhiteList;
                }
                "[accept_all]" | "[proxy_all]" => {
                    mode = Mode::BlackList;
                }
                "[outbound_block_list]" => {
                    curr_ipv4 = &mut outbound_block_ipv4;
                    curr_ipv6 = &mut outbound_block_ipv6;
                    curr_rules = &mut outbound_block_rules;
                }
                "[black_list]" | "[bypass_list]" => {
                    curr_ipv4 = &mut bypass_ipv4;
                    curr_ipv6 = &mut bypass_ipv6;
                    curr_rules = &mut bypass_rules;
                }
                "[white_list]" | "[proxy_list]" => {
                    curr_ipv4 = &mut proxy_ipv4;
                    curr_ipv6 = &mut proxy_ipv6;
                    curr_rules = &mut proxy_rules;
                }
                _ => {
                    match line.parse::<IpNet>() {
                        Ok(IpNet::V4(v4)) => {
                            curr_ipv4.add(v4);
                        }
                        Ok(IpNet::V6(v6)) => {
                            curr_ipv6.add(v6);
                        }
                        Err(..) => {
                            // Maybe it is a pure IpAddr
                            match line.parse::<IpAddr>() {
                                Ok(IpAddr::V4(v4)) => {
                                    curr_ipv4.add(Ipv4Net::from(v4));
                                }
                                Ok(IpAddr::V6(v6)) => {
                                    curr_ipv6.add(Ipv6Net::from(v6));
                                }
                                Err(..) => {
                                    // FIXME: If this line is not a valid regex, how can we know without actually compile it?
                                    curr_rules.push(line);
                                }
                            }
                        }
                    }
                }
            }
        }

        const REGEX_SIZE_LIMIT: usize = usize::max_value();

        let outbound_block_regex = match RegexSetBuilder::new(outbound_block_rules)
            .size_limit(REGEX_SIZE_LIMIT)
            .build()
        {
            Ok(r) => r,
            Err(err) => {
                let err = Error::new(ErrorKind::Other, format!("[outbound_block_list] regex error: {}", err));
                return Err(err);
            }
        };

        let bypass_regex = match RegexSetBuilder::new(bypass_rules).size_limit(REGEX_SIZE_LIMIT).build() {
            Ok(r) => r,
            Err(err) => {
                let err = Error::new(
                    ErrorKind::Other,
                    format!("[black_list] or [bypass_list] regex error: {}", err),
                );
                return Err(err);
            }
        };

        let proxy_regex = match RegexSetBuilder::new(proxy_rules).size_limit(REGEX_SIZE_LIMIT).build() {
            Ok(r) => r,
            Err(err) => {
                let err = Error::new(
                    ErrorKind::Other,
                    format!("[white_list] or [proxy_list] regex error: {}", err),
                );
                return Err(err);
            }
        };

        Ok(AccessControl {
            outbound_block: Rules::new(outbound_block_ipv4, outbound_block_ipv6, outbound_block_regex),
            black_list: Rules::new(bypass_ipv4, bypass_ipv6, bypass_regex),
            white_list: Rules::new(proxy_ipv4, proxy_ipv6, proxy_regex),
            mode,
        })
    }

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    pub fn check_query_in_proxy_list(&self, query: &Query) -> Option<bool> {
        if query.query_class() != DNSClass::IN {
            // unconditionally forward all non-IN queries to remote
            return Some(true);
        }
        // remove the last dot from fqdn name
        let mut name = query.name().to_ascii();
        name.pop();
        let addr = Address::DomainNameAddress(name, 0);
        // Addresses in proxy_list will be proxied
        if self.white_list.check_address_matched(&addr) {
            return Some(true);
        }
        if self.black_list.check_address_matched(&addr) {
            return Some(false);
        }
        match self.mode {
            Mode::BlackList => if self.black_list.is_rule_empty_for_qtype(query.query_type()) {
                return Some(true);
            },
            Mode::WhiteList => if self.white_list.is_rule_empty_for_qtype(query.query_type()) {
                return Some(false);
            },
        }
        None
    }

    pub fn check_ip_in_proxy_list(&self, ip: &IpAddr) -> bool {
        match self.mode {
            Mode::BlackList => !self.black_list.check_ip_matched(ip),
            Mode::WhiteList => self.white_list.check_ip_matched(ip),
        }
    }

    /// Check if target address should be bypassed (for client)
    ///
    /// FIXME: This function may perform a DNS resolution
    pub async fn check_target_bypassed(&self, context: &Context, addr: &Address) -> bool {
        // Addresses in bypass_list will be bypassed
        if self.black_list.check_address_matched(addr) {
            return true;
        }

        // Addresses in proxy_list will be proxied
        if self.white_list.check_address_matched(addr) {
            return false;
        }

        // Resolve hostname and check the list
        if cfg!(not(target_os = "android")) {
            if let Address::DomainNameAddress(ref host, port) = *addr {
                if let Ok(vaddr) = context.dns_resolve(host, port).await {
                    for addr in vaddr {
                        if self.black_list.check_ip_matched(&addr.ip()) {
                            return true;
                        }

                        if self.white_list.check_ip_matched(&addr.ip()) {
                            return false;
                        }
                    }
                }
            }
        }

        // default rule
        match self.mode {
            Mode::BlackList => false,
            Mode::WhiteList => true,
        }
    }

    /// Check if client address should be blocked (for server)
    pub fn check_client_blocked(&self, addr: &SocketAddr) -> bool {
        match self.mode {
            Mode::BlackList => {
                // Only clients in black_list will be blocked
                self.black_list.check_ip_matched(&addr.ip())
            }
            Mode::WhiteList => {
                // Only clients in white_list will be proxied
                !self.white_list.check_ip_matched(&addr.ip())
            }
        }
    }

    /// Check if outbound address is blocked (for server)
    ///
    /// NOTE: `Address::DomainName` is only validated by regex rules,
    ///       resolved addresses are checked in the `lookup_outbound_then!` macro
    pub fn check_outbound_blocked(&self, outbound: &Address) -> bool {
        self.outbound_block.check_address_matched(outbound)
    }

    /// Check resolved outbound address is blocked (for server)
    pub fn check_resolved_outbound_blocked(&self, outbound: &SocketAddr) -> bool {
        self.outbound_block.check_ip_matched(&outbound.ip())
    }
}
<|endoftext|>