<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::{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())
    }
}
<commit_msg>Respect default mode for fallback<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 use default for all non-IN queries
            return Some(self.is_default_in_proxy_list());
        }
        // 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),
        }
    }

    fn is_default_in_proxy_list(&self) -> bool {
        match self.mode {
            Mode::BlackList => true,
            Mode::WhiteList => false,
        }
    }

    /// 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;
                        }
                    }
                }
            }
        }

        !self.is_default_in_proxy_list()
    }

    /// 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|>


<commit_before>use std::cell::{RefCell, Ref};

use opencl;
use opencl::hl::KernelArg;
use opencl::mem::{Buffer, CLBuffer};
use libc;
use ref_filter_map::ref_filter_map;

use array::Array;
use context::Context;
use helper;
use num::Num;
use range_arg::RangeArg;

pub enum TensorMode {
    In,
    Out,
    Mut,
}

pub struct Tensor<T: Num> {
    shape: Vec<usize>,
    dim_steps: Vec<usize>,
    buffer: CLBuffer<T>,
    event: RefCell<Option<Event>>,
}

impl<T: Num> Tensor<T> {
    pub fn new(ctx: &Context, shape: Vec<usize>, mode: TensorMode) -> Tensor<T> {
        let mem_mode =
            match mode {
                TensorMode::In => { opencl::cl::CL_MEM_READ_ONLY },
                TensorMode::Out => { opencl::cl::CL_MEM_WRITE_ONLY },
                TensorMode::Mut => { opencl::cl::CL_MEM_READ_WRITE },
            };
        let buf_size = shape.iter().fold(1, |a, b| a*b);
        let dim_steps = helper::compute_dim_steps(&shape);
        Tensor {
            shape: shape,
            dim_steps: dim_steps,
            buffer: ctx.ctx.create_buffer(buf_size, mem_mode),
            event: RefCell::new(None),
        }
    }

    pub fn from_array(ctx: &Context,
                      array: &Array<T>,
                      mode: TensorMode) -> Tensor<T> {
        let mem_mode =
            match mode {
                TensorMode::In => { opencl::cl::CL_MEM_READ_ONLY },
                TensorMode::Out => { opencl::cl::CL_MEM_WRITE_ONLY },
                TensorMode::Mut => { opencl::cl::CL_MEM_READ_WRITE },
            };
        Tensor {
            shape: array.shape().to_vec(),
            dim_steps: array.dim_steps().to_owned(),
            buffer: ctx.ctx.create_buffer_from(array.buffer(), mem_mode),
            event: RefCell::new(None),
        }
    }

    pub fn get(&self, ctx: &Context) -> Array<T> {
        let vec = ctx.queue.get(&self.buffer, self.get_event().as_ref().map(|x| &**x));
        Array::from_vec(self.shape.clone(), vec)
    }

    pub fn read(&self, ctx: &Context, array: &mut Array<T>) {
        ctx.queue.read(&self.buffer, &mut array.buffer_mut(), self.get_event().as_ref().map(|x| &**x));
    }

    pub fn set(&self, ctx: &Context, array: &Array<T>) {
        ctx.queue.write(&self.buffer, &array.buffer(), ());
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    pub fn set_event(&self, e: Event) {
        *self.event.borrow_mut() = Some(e);
    }

    pub fn get_event(&self) -> Option<Ref<Event>> {
        if self.event.borrow().is_some() {
            ref_filter_map(self.event.borrow(), |o| o.as_ref())
        } else {
            None
        }
    }

    pub fn slice<'t, 'r>(&'t self, r: &'r [RangeArg]) -> TensorView<'t, 'r, T> {
        TensorView {
            shape: self.shape.as_ref(),
            dim_steps: self.dim_steps.as_ref(),
            ranges: r,
            buffer: &self.buffer,
            event: &self.event,
        }
    }
}

impl<T: Num> KernelArg for Tensor<T> {
    fn get_value(&self) -> (libc::size_t, *const libc::c_void) {
        self.buffer.get_value()
    }
}

impl<'t, 'r, T: Num> KernelArg for TensorView<'t, 'r, T> {
    fn get_value(&self) -> (libc::size_t, *const libc::c_void) {
        self.buffer.get_value()
    }
}

pub struct TensorView<'t, 'r, T: Num+'t> {
    pub shape: &'t [usize],
    pub dim_steps: &'t [usize],
    ranges: &'r [RangeArg],
    buffer: &'t CLBuffer<T>,
    event: &'t RefCell<Option<Event>>,
}

impl<'t, 'r, T: Num> TensorView<'t, 'r, T> {
    pub fn set_event(&self, e: Event) {
        *self.event.borrow_mut() = Some(e);
    }

    pub fn get_event(&self) -> Option<Ref<Event>> {
        if self.event.borrow().is_some() {
            ref_filter_map(self.event.borrow(), |o| o.as_ref())
        } else {
            None
        }
    }

    pub fn view_offset(&self, dim: usize) -> usize {
        self.ranges[dim].start
    }

    pub fn view_shape(&self, dim: usize) -> usize {
        self.ranges[dim].len(self.shape[dim])
    }

    pub fn len(&self) -> usize {
        self.buffer.len()
    }
}

pub type Event = opencl::hl::Event;

////////////////////////////////////////////////////////////////////////////////////////////////////

#[test]
fn test_tensor_slicing() {
    use ops::add_slice;

    let ref ctx = Context::new();

    let a = Array::from_vec(vec![4, 3], vec![2, 3, 4,
                                             6, 7, 8,
                                             10, 11, 12,
                                             14, 15, 16]);
    let at = Tensor::from_array(ctx, &a, TensorMode::Mut);
    let slice: &[RangeArg] = &[RangeArg::from(1..3), RangeArg::from(1)];
    let atv = at.slice(slice);

    let b = Array::from_vec(vec![4, 4], vec![1, 2, 3, 4,
                                             5, 6, 7, 8,
                                             9, 10, 11, 12,
                                             13, 14, 15, 16]);
    let bt = Tensor::from_array(ctx, &b, TensorMode::Mut);
    let slice: &[RangeArg] = &[RangeArg::from(1..3), RangeArg::from(3)];
    let btv = bt.slice(slice);

    let c = Array::from_vec(vec![4, 4], vec![0; 16]);
    let ct = Tensor::from_array(ctx, &c, TensorMode::Mut);
    let slice: &[RangeArg] = &[RangeArg::from(2..4), RangeArg::from(0)];
    let ctv = ct.slice(slice);

    add_slice(ctx, &atv, &btv, &ctv);
    assert!(ct.get(ctx).buffer() == &[0, 0, 0, 0,
                                      0, 0, 0, 0,
                                      15, 0, 0, 0,
                                      23, 0, 0, 0]);
}
<commit_msg>Added unit test for Tensor::read<commit_after>use std::cell::{RefCell, Ref};

use opencl;
use opencl::hl::KernelArg;
use opencl::mem::{Buffer, CLBuffer};
use libc;
use ref_filter_map::ref_filter_map;

use array::Array;
use context::Context;
use helper;
use num::Num;
use range_arg::RangeArg;

pub enum TensorMode {
    In,
    Out,
    Mut,
}

pub struct Tensor<T: Num> {
    shape: Vec<usize>,
    dim_steps: Vec<usize>,
    buffer: CLBuffer<T>,
    event: RefCell<Option<Event>>,
}

impl<T: Num> Tensor<T> {
    pub fn new(ctx: &Context, shape: Vec<usize>, mode: TensorMode) -> Tensor<T> {
        let mem_mode =
            match mode {
                TensorMode::In => { opencl::cl::CL_MEM_READ_ONLY },
                TensorMode::Out => { opencl::cl::CL_MEM_WRITE_ONLY },
                TensorMode::Mut => { opencl::cl::CL_MEM_READ_WRITE },
            };
        let buf_size = shape.iter().fold(1, |a, b| a*b);
        let dim_steps = helper::compute_dim_steps(&shape);
        Tensor {
            shape: shape,
            dim_steps: dim_steps,
            buffer: ctx.ctx.create_buffer(buf_size, mem_mode),
            event: RefCell::new(None),
        }
    }

    pub fn from_array(ctx: &Context,
                      array: &Array<T>,
                      mode: TensorMode) -> Tensor<T> {
        let mem_mode =
            match mode {
                TensorMode::In => { opencl::cl::CL_MEM_READ_ONLY },
                TensorMode::Out => { opencl::cl::CL_MEM_WRITE_ONLY },
                TensorMode::Mut => { opencl::cl::CL_MEM_READ_WRITE },
            };
        Tensor {
            shape: array.shape().to_vec(),
            dim_steps: array.dim_steps().to_owned(),
            buffer: ctx.ctx.create_buffer_from(array.buffer(), mem_mode),
            event: RefCell::new(None),
        }
    }

    pub fn get(&self, ctx: &Context) -> Array<T> {
        let vec = ctx.queue.get(&self.buffer, self.get_event().as_ref().map(|x| &**x));
        Array::from_vec(self.shape.clone(), vec)
    }

    pub fn read(&self, ctx: &Context, array: &mut Array<T>) {
        ctx.queue.read(&self.buffer, &mut array.buffer_mut(), self.get_event().as_ref().map(|x| &**x));
    }

    pub fn set(&self, ctx: &Context, array: &Array<T>) {
        ctx.queue.write(&self.buffer, &array.buffer(), ());
    }

    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    pub fn set_event(&self, e: Event) {
        *self.event.borrow_mut() = Some(e);
    }

    pub fn get_event(&self) -> Option<Ref<Event>> {
        if self.event.borrow().is_some() {
            ref_filter_map(self.event.borrow(), |o| o.as_ref())
        } else {
            None
        }
    }

    pub fn slice<'t, 'r>(&'t self, r: &'r [RangeArg]) -> TensorView<'t, 'r, T> {
        TensorView {
            shape: self.shape.as_ref(),
            dim_steps: self.dim_steps.as_ref(),
            ranges: r,
            buffer: &self.buffer,
            event: &self.event,
        }
    }
}

impl<T: Num> KernelArg for Tensor<T> {
    fn get_value(&self) -> (libc::size_t, *const libc::c_void) {
        self.buffer.get_value()
    }
}

impl<'t, 'r, T: Num> KernelArg for TensorView<'t, 'r, T> {
    fn get_value(&self) -> (libc::size_t, *const libc::c_void) {
        self.buffer.get_value()
    }
}

pub struct TensorView<'t, 'r, T: Num+'t> {
    pub shape: &'t [usize],
    pub dim_steps: &'t [usize],
    ranges: &'r [RangeArg],
    buffer: &'t CLBuffer<T>,
    event: &'t RefCell<Option<Event>>,
}

impl<'t, 'r, T: Num> TensorView<'t, 'r, T> {
    pub fn set_event(&self, e: Event) {
        *self.event.borrow_mut() = Some(e);
    }

    pub fn get_event(&self) -> Option<Ref<Event>> {
        if self.event.borrow().is_some() {
            ref_filter_map(self.event.borrow(), |o| o.as_ref())
        } else {
            None
        }
    }

    pub fn view_offset(&self, dim: usize) -> usize {
        self.ranges[dim].start
    }

    pub fn view_shape(&self, dim: usize) -> usize {
        self.ranges[dim].len(self.shape[dim])
    }

    pub fn len(&self) -> usize {
        self.buffer.len()
    }
}

pub type Event = opencl::hl::Event;

////////////////////////////////////////////////////////////////////////////////////////////////////

#[test]
fn test_tensor_read() {
    use ops::add_slice;

    let ref ctx = Context::new();

    let a = Array::from_vec(vec![2, 2], vec![1i32, 2,
                                             3, 4]);
    let at = Tensor::from_array(ctx, &a, TensorMode::Out);
    let mut b = Array::new(vec![2, 2], 0);

    at.read(ctx, &mut b);

    assert!(b.buffer() == &[1, 2,
                            3, 4]);
}

#[test]
fn test_tensor_slicing() {
    use ops::add_slice;

    let ref ctx = Context::new();

    let a = Array::from_vec(vec![4, 3], vec![2, 3, 4,
                                             6, 7, 8,
                                             10, 11, 12,
                                             14, 15, 16]);
    let at = Tensor::from_array(ctx, &a, TensorMode::Mut);
    let slice: &[RangeArg] = &[RangeArg::from(1..3), RangeArg::from(1)];
    let atv = at.slice(slice);

    let b = Array::from_vec(vec![4, 4], vec![1, 2, 3, 4,
                                             5, 6, 7, 8,
                                             9, 10, 11, 12,
                                             13, 14, 15, 16]);
    let bt = Tensor::from_array(ctx, &b, TensorMode::Mut);
    let slice: &[RangeArg] = &[RangeArg::from(1..3), RangeArg::from(3)];
    let btv = bt.slice(slice);

    let c = Array::from_vec(vec![4, 4], vec![0; 16]);
    let ct = Tensor::from_array(ctx, &c, TensorMode::Mut);
    let slice: &[RangeArg] = &[RangeArg::from(2..4), RangeArg::from(0)];
    let ctv = ct.slice(slice);

    add_slice(ctx, &atv, &btv, &ctv);
    assert!(ct.get(ctx).buffer() == &[0, 0, 0, 0,
                                      0, 0, 0, 0,
                                      15, 0, 0, 0,
                                      23, 0, 0, 0]);
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use super::model;
use regex;
use rs_es;
use rs_es::query::Query as rs_q;
use rs_es::operations::search::SearchResult;
use mimir;
use serde_json;

fn build_rs_client(cnx: &String) -> rs_es::Client {
    let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)/(?P<index>\w+)")
                 .unwrap();
    let cap = re.captures(&cnx).unwrap();
    let host = cap.name("host").unwrap();
    let port = cap.name("port").unwrap().parse::<u32>().unwrap();

    rs_es::Client::new(&host, port)
}

/// takes a ES json blob and build a Place from it
/// it uses the _type field of ES to know which type of the Place enum to fill
fn make_place(doc_type: String, value: Option<Box<serde_json::Value>>) -> Option<mimir::Place> {
    value.and_then(|v| {
        match doc_type.as_ref() {
            "addr" => serde_json::from_value::<mimir::Addr>(*v).ok().and_then(|o| Some(mimir::Place::Addr(o))),
            "street" => serde_json::from_value::<mimir::Street>(*v).ok().and_then(|o| Some(mimir::Place::Street(o))),
            "admin" => serde_json::from_value::<mimir::Admin>(*v).ok().and_then(|o| Some(mimir::Place::Admin(o))),
            _    => {
                warn!("unknown ES return value, _type field = {}", doc_type);
                None
            }
        }
    })
}

fn query(q: &String, cnx: &String, match_type: &str) -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    let sub_query = rs_q::build_bool()
                        .with_should(vec![
                       rs_q::build_term("_type","addr").with_boost(1000).build(),
                       rs_q::build_match(match_type, q.to_string())
                              .with_boost(100)
                              .build(),
                       rs_q::build_function_score()
                              .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
                              .with_boost(30)
                              .with_query(rs_q::build_match_all().build())
			                  .with_function(
                          rs_es::query::functions::Function::build_field_value_factor("weight")
			                                  .with_factor(1)
                                      .with_modifier(rs_es::query::functions::Modifier::Log1p)
			                                                          .build())
                              .build()])
                        .build();
    let filter = rs_q::build_bool()
                     .with_should(vec![rs_q::build_bool()
                                           .with_must_not(rs_q::build_exists("house_number")
                                                              .build())
                                           .build(),
                                       rs_q::build_match("house_number", q.to_string()).build()])
                     .with_must(vec![rs_q::build_match(match_type, q.to_string())
             .with_minimum_should_match(rs_es::query::MinimumShouldMatch::from(100f64)).build()])
                     .build();

    let final_query = rs_q::build_bool()
                          .with_must(vec![sub_query])
                          .with_filter(filter)
                          .build();

    let mut client = build_rs_client(cnx);

    let result: SearchResult<serde_json::Value> = try!(client.search_query()
                                                       .with_indexes(&["munin"])
                                                       .with_query(&final_query)
                                                       .send());

    debug!("{} documents found", result.hits.total);

    // for the moment rs-es does not handle enum Document,
    // so we need to convert the ES glob to a Place
    Ok(result.hits
             .hits
             .into_iter()
             .filter_map(|hit| make_place(hit.doc_type, hit.source))
             .collect())
}

fn query_location(_q: &String,
                  _coord: &model::Coord)
                  -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    panic!("todo!");
}

fn query_prefix(q: &String, cnx: &String) -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
	query(&q, cnx, "name.prefix")
}

fn query_ngram(q: &String, cnx: &String) -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
	query(&q, cnx, "name.ngram")
}

pub fn autocomplete(q: String,
                    coord: Option<model::Coord>,
                    cnx: &String)
                    -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    if let Some(ref coord) = coord {
        query_location(&q, coord)
    } else {
    	//First search with match = "name.prefix".
    	//If no result then another search with match = "name.ngram"
    	let results = try!(query_prefix(&q, cnx));
        if results.is_empty() {
        	query_ngram(&q, cnx)
        } else {
        	Ok(results)
        }
    }
}
<commit_msg>fix es url<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use super::model;
use regex;
use rs_es;
use rs_es::query::Query as rs_q;
use rs_es::operations::search::SearchResult;
use mimir;
use serde_json;

fn build_rs_client(cnx: &String) -> rs_es::Client {
    let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)")
                 .unwrap();
    let cap = re.captures(&cnx).unwrap();
    let host = cap.name("host").unwrap();
    let port = cap.name("port").unwrap().parse::<u32>().unwrap();

    rs_es::Client::new(&host, port)
}

/// takes a ES json blob and build a Place from it
/// it uses the _type field of ES to know which type of the Place enum to fill
fn make_place(doc_type: String, value: Option<Box<serde_json::Value>>) -> Option<mimir::Place> {
    value.and_then(|v| {
        match doc_type.as_ref() {
            "addr" => serde_json::from_value::<mimir::Addr>(*v).ok().and_then(|o| Some(mimir::Place::Addr(o))),
            "street" => serde_json::from_value::<mimir::Street>(*v).ok().and_then(|o| Some(mimir::Place::Street(o))),
            "admin" => serde_json::from_value::<mimir::Admin>(*v).ok().and_then(|o| Some(mimir::Place::Admin(o))),
            _    => {
                warn!("unknown ES return value, _type field = {}", doc_type);
                None
            }
        }
    })
}

fn query(q: &String, cnx: &String, match_type: &str) -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    let sub_query = rs_q::build_bool()
                        .with_should(vec![
                       rs_q::build_term("_type","addr").with_boost(1000).build(),
                       rs_q::build_match(match_type, q.to_string())
                              .with_boost(100)
                              .build(),
                       rs_q::build_function_score()
                              .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
                              .with_boost(30)
                              .with_query(rs_q::build_match_all().build())
			                  .with_function(
                          rs_es::query::functions::Function::build_field_value_factor("weight")
			                                  .with_factor(1)
                                      .with_modifier(rs_es::query::functions::Modifier::Log1p)
			                                                          .build())
                              .build()])
                        .build();
    let filter = rs_q::build_bool()
                     .with_should(vec![rs_q::build_bool()
                                           .with_must_not(rs_q::build_exists("house_number")
                                                              .build())
                                           .build(),
                                       rs_q::build_match("house_number", q.to_string()).build()])
                     .with_must(vec![rs_q::build_match(match_type, q.to_string())
             .with_minimum_should_match(rs_es::query::MinimumShouldMatch::from(100f64)).build()])
                     .build();

    let final_query = rs_q::build_bool()
                          .with_must(vec![sub_query])
                          .with_filter(filter)
                          .build();

    let mut client = build_rs_client(cnx);

    let result: SearchResult<serde_json::Value> = try!(client.search_query()
                                                       .with_indexes(&["munin"])
                                                       .with_query(&final_query)
                                                       .send());

    debug!("{} documents found", result.hits.total);

    // for the moment rs-es does not handle enum Document,
    // so we need to convert the ES glob to a Place
    Ok(result.hits
             .hits
             .into_iter()
             .filter_map(|hit| make_place(hit.doc_type, hit.source))
             .collect())
}

fn query_location(_q: &String,
                  _coord: &model::Coord)
                  -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    panic!("todo!");
}

fn query_prefix(q: &String, cnx: &String) -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
	query(&q, cnx, "name.prefix")
}

fn query_ngram(q: &String, cnx: &String) -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
	query(&q, cnx, "name.ngram")
}

pub fn autocomplete(q: String,
                    coord: Option<model::Coord>,
                    cnx: &String)
                    -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    if let Some(ref coord) = coord {
        query_location(&q, coord)
    } else {
    	//First search with match = "name.prefix".
    	//If no result then another search with match = "name.ngram"
    	let results = try!(query_prefix(&q, cnx));
        if results.is_empty() {
        	query_ngram(&q, cnx)
        } else {
        	Ok(results)
        }
    }
}
<|endoftext|>


<commit_before>use mio::*;
use mio::deprecated::{Handler, EventLoop};
use mio::net::UdpSocket;
use bytes::{Buf, RingBuf, SliceBuf, MutBuf};
use std::io::ErrorKind;
use std::str;
use std::time;
use localhost;

const LISTENER: Token = Token(0);
const SENDER: Token = Token(1);

pub struct UdpHandlerSendRecv {
    tx: UdpSocket,
    rx: UdpSocket,
    msg: &'static str,
    buf: SliceBuf<'static>,
    rx_buf: RingBuf,
    connected: bool
}

impl UdpHandlerSendRecv {
    fn new(tx: UdpSocket, rx: UdpSocket, connected: bool, msg : &'static str) -> UdpHandlerSendRecv {
        UdpHandlerSendRecv {
            tx: tx,
            rx: rx,
            msg: msg,
            buf: SliceBuf::wrap(msg.as_bytes()),
            rx_buf: RingBuf::new(1024),
            connected: connected
        }
    }
}

impl Handler for UdpHandlerSendRecv {
    type Timeout = usize;
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<UdpHandlerSendRecv>, token: Token, events: Ready) {

        if events.is_readable() {
            match token {
                LISTENER => {
                    debug!("We are receiving a datagram now...");
                    let cnt = unsafe {
                        if !self.connected {
                            self.rx.recv_from(self.rx_buf.mut_bytes()).unwrap().0
                        } else {
                            self.rx.recv(self.rx_buf.mut_bytes()).unwrap()
                        }
                    };

                    unsafe { MutBuf::advance(&mut self.rx_buf, cnt); }
                    assert!(str::from_utf8(self.rx_buf.bytes()).unwrap() == self.msg);
                    event_loop.shutdown();
                },
                _ => ()
            }
        }

        if events.is_writable() {
            match token {
                SENDER => {
                    let cnt = if !self.connected {
                        let addr = self.rx.local_addr().unwrap();
                        self.tx.send_to(self.buf.bytes(), &addr).unwrap()
                    } else {
                        self.tx.send(self.buf.bytes()).unwrap()
                    };

                    self.buf.advance(cnt);
                },
                _ => {}
            }
        }
    }
}

pub struct UdpHandlerTimeout {
}

impl Handler for UdpHandlerTimeout {
    type Timeout = usize;
    type Message = ();

    fn ready(&mut self, _event_loop: &mut EventLoop<UdpHandlerTimeout>, token: Token, events: Ready) {
        if events.is_readable() {
            match token {
                LISTENER => {
                    assert!(false, "Expected to no receive a packet but got something")
                },
                _ => ()
            }
        }
   }

    /// Invoked when a timeout has completed.
    fn timeout(&mut self, event_loop: &mut EventLoop<Self>, _timeout: Self::Timeout) {
        event_loop.shutdown();
    }
}

fn assert_send<T: Send>() {
}

fn assert_sync<T: Sync>() {
}

#[cfg(test)]
fn test_send_recv_udp(tx: UdpSocket, rx: UdpSocket, connected: bool) {
    debug!("Starting TEST_UDP_SOCKETS");
    let mut event_loop = EventLoop::new().unwrap();

    assert_send::<UdpSocket>();
    assert_sync::<UdpSocket>();

    // ensure that the sockets are non-blocking
    let mut buf = [0; 128];
    assert_eq!(ErrorKind::WouldBlock, rx.recv_from(&mut buf).unwrap_err().kind());

    info!("Registering SENDER");
    event_loop.register(&tx, SENDER, Ready::writable(), PollOpt::edge()).unwrap();

    info!("Registering LISTENER");
    event_loop.register(&rx, LISTENER, Ready::readable(), PollOpt::edge()).unwrap();

    info!("Starting event loop to test with...");
    event_loop.run(&mut UdpHandlerSendRecv::new(tx, rx, connected, "hello world")).unwrap();
}

#[test]
pub fn test_udp_socket() {
    let addr = localhost();
    let any = localhost();

    let tx = UdpSocket::bind(&any).unwrap();
    let rx = UdpSocket::bind(&addr).unwrap();

    test_send_recv_udp(tx, rx, false);
}

#[test]
pub fn test_udp_socket_send_recv() {
    let addr = localhost();
    let any = localhost();

    let tx = UdpSocket::bind(&any).unwrap();
    let rx = UdpSocket::bind(&addr).unwrap();

    let tx_addr = tx.local_addr().unwrap();
    let rx_addr = rx.local_addr().unwrap();

    assert!(tx.connect(rx_addr).is_ok());
    assert!(rx.connect(tx_addr).is_ok());

    test_send_recv_udp(tx, rx, true);
}

#[test]
pub fn test_udp_socket_discard() {
    let addr = localhost();
    let any = localhost();
    let outside = localhost();

    let tx = UdpSocket::bind(&any).unwrap();
    let rx = UdpSocket::bind(&addr).unwrap();
    let udp_outside = UdpSocket::bind(&outside).unwrap();

    let tx_addr = tx.local_addr().unwrap();
    let rx_addr = rx.local_addr().unwrap();

    assert!(tx.connect(rx_addr).is_ok());
    assert!(udp_outside.connect(rx_addr).is_ok());
    assert!(rx.connect(tx_addr).is_ok());

    let mut event_loop = EventLoop::new().unwrap();

    let r = udp_outside.send("hello world".as_bytes());
    assert!(r.is_ok() || r.unwrap_err().kind() == ErrorKind::WouldBlock);

    event_loop.register(&rx, LISTENER, Ready::readable(), PollOpt::edge()).unwrap();
    event_loop.register(&tx, SENDER, Ready::writable(), PollOpt::edge()).unwrap();

    event_loop.timeout(5000, time::Duration::from_secs(5)).unwrap();
    event_loop.run(&mut UdpHandlerTimeout {}).unwrap();
}
<commit_msg>Change test_udp_socket to use Poll API. (#716)<commit_after>use mio::{Events, Poll, PollOpt, Ready, Token};
use mio::net::UdpSocket;
use bytes::{Buf, RingBuf, SliceBuf, MutBuf};
use std::io::ErrorKind;
use std::str;
use std::time;
use localhost;

const LISTENER: Token = Token(0);
const SENDER: Token = Token(1);

pub struct UdpHandlerSendRecv {
    tx: UdpSocket,
    rx: UdpSocket,
    msg: &'static str,
    buf: SliceBuf<'static>,
    rx_buf: RingBuf,
    connected: bool,
    shutdown: bool,
}

impl UdpHandlerSendRecv {
    fn new(tx: UdpSocket, rx: UdpSocket, connected: bool, msg : &'static str) -> UdpHandlerSendRecv {
        UdpHandlerSendRecv {
            tx: tx,
            rx: rx,
            msg: msg,
            buf: SliceBuf::wrap(msg.as_bytes()),
            rx_buf: RingBuf::new(1024),
            connected: connected,
            shutdown: false,
        }
    }
}

fn assert_send<T: Send>() {
}

fn assert_sync<T: Sync>() {
}

#[cfg(test)]
fn test_send_recv_udp(tx: UdpSocket, rx: UdpSocket, connected: bool) {
    debug!("Starting TEST_UDP_SOCKETS");
    let poll = Poll::new().unwrap();

    assert_send::<UdpSocket>();
    assert_sync::<UdpSocket>();

    // ensure that the sockets are non-blocking
    let mut buf = [0; 128];
    assert_eq!(ErrorKind::WouldBlock, rx.recv_from(&mut buf).unwrap_err().kind());

    info!("Registering SENDER");
    poll.register(&tx, SENDER, Ready::writable(), PollOpt::edge()).unwrap();

    info!("Registering LISTENER");
    poll.register(&rx, LISTENER, Ready::readable(), PollOpt::edge()).unwrap();

    let mut events = Events::with_capacity(1024);

    info!("Starting event loop to test with...");
    let mut handler = UdpHandlerSendRecv::new(tx, rx, connected, "hello world");

    while !handler.shutdown {
        poll.poll(&mut events, None).unwrap();

        for event in &events {
            if event.readiness().is_readable() {
                match event.token() {
                    LISTENER => {
                        debug!("We are receiving a datagram now...");
                        let cnt = unsafe {
                            if !handler.connected {
                                handler.rx.recv_from(handler.rx_buf.mut_bytes()).unwrap().0
                            } else {
                                handler.rx.recv(handler.rx_buf.mut_bytes()).unwrap()
                            }
                        };

                        unsafe { MutBuf::advance(&mut handler.rx_buf, cnt); }
                        assert!(str::from_utf8(handler.rx_buf.bytes()).unwrap() == handler.msg);
                        handler.shutdown = true;
                    },
                    _ => ()
                }
            }

            if event.readiness().is_writable() {
                match event.token() {
                    SENDER => {
                        let cnt = if !handler.connected {
                            let addr = handler.rx.local_addr().unwrap();
                            handler.tx.send_to(handler.buf.bytes(), &addr).unwrap()
                        } else {
                            handler.tx.send(handler.buf.bytes()).unwrap()
                        };

                        handler.buf.advance(cnt);
                    },
                    _ => {}
                }
            }
        }
    }
}

#[test]
pub fn test_udp_socket() {
    let addr = localhost();
    let any = localhost();

    let tx = UdpSocket::bind(&any).unwrap();
    let rx = UdpSocket::bind(&addr).unwrap();

    test_send_recv_udp(tx, rx, false);
}

#[test]
pub fn test_udp_socket_send_recv() {
    let addr = localhost();
    let any = localhost();

    let tx = UdpSocket::bind(&any).unwrap();
    let rx = UdpSocket::bind(&addr).unwrap();

    let tx_addr = tx.local_addr().unwrap();
    let rx_addr = rx.local_addr().unwrap();

    assert!(tx.connect(rx_addr).is_ok());
    assert!(rx.connect(tx_addr).is_ok());

    test_send_recv_udp(tx, rx, true);
}

#[test]
pub fn test_udp_socket_discard() {
    let addr = localhost();
    let any = localhost();
    let outside = localhost();

    let tx = UdpSocket::bind(&any).unwrap();
    let rx = UdpSocket::bind(&addr).unwrap();
    let udp_outside = UdpSocket::bind(&outside).unwrap();

    let tx_addr = tx.local_addr().unwrap();
    let rx_addr = rx.local_addr().unwrap();

    assert!(tx.connect(rx_addr).is_ok());
    assert!(udp_outside.connect(rx_addr).is_ok());
    assert!(rx.connect(tx_addr).is_ok());

    let poll = Poll::new().unwrap();

    let r = udp_outside.send("hello world".as_bytes());
    assert!(r.is_ok() || r.unwrap_err().kind() == ErrorKind::WouldBlock);

    poll.register(&rx, LISTENER, Ready::readable(), PollOpt::edge()).unwrap();
    poll.register(&tx, SENDER, Ready::writable(), PollOpt::edge()).unwrap();

    let mut events = Events::with_capacity(1024);

    poll.poll(&mut events, Some(time::Duration::from_secs(5))).unwrap();

    for event in &events {
        if event.readiness().is_readable() {
            match event.token() {
                LISTENER => {
                    assert!(false, "Expected to no receive a packet but got something")
                },
                _ => ()
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    fs,
    thread,
    path::{Path, PathBuf},
    sync::{mpsc, Arc},
    time::Duration,
};
use crossbeam_channel::{Receiver, Sender, unbounded, RecvError, select};
use relative_path::RelativePathBuf;
use thread_worker::WorkerHandle;
use walkdir::WalkDir;
use notify::{DebouncedEvent, RecommendedWatcher, RecursiveMode, Watcher as _Watcher};

use crate::{RootConfig, Roots, VfsRoot};

type Result<T> = std::result::Result<T, crossbeam_channel::SendError<TaskResult>>;

pub(crate) enum Task {
    AddRoot {
        root: VfsRoot,
        config: Arc<RootConfig>,
    },
}

#[derive(Debug)]
pub enum TaskResult {
    BulkLoadRoot {
        root: VfsRoot,
        files: Vec<(RelativePathBuf, String)>,
    },
    AddSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
    },
    ChangeSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
    },
    RemoveSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
    },
}

#[derive(Debug)]
enum ChangeKind {
    Create,
    Write,
    Remove,
}

const WATCHER_DELAY: Duration = Duration::from_millis(250);

pub(crate) struct Worker {
    worker: thread_worker::Worker<Task, TaskResult>,
    worker_handle: WorkerHandle,
}

impl Worker {
    pub(crate) fn start(roots: Arc<Roots>) -> Worker {
        // This is a pretty elaborate setup of threads & channels! It is
        // explained by the following concerns:

        //    * we need to burn a thread translating from notify's mpsc to
        //      crossbeam_channel.
        //    * we want to read all files from a single thread, to gurantee that
        //      we always get fresher versions and never go back in time.
        //    * we want to tear down everything neatly during shutdown.
        let (worker, worker_handle) = thread_worker::spawn(
            "vfs",
            128,
            // This are the channels we use to communicate with outside world.
            // If `input_receiver` is closed we need to tear ourselves down.
            // `output_sender` should not be closed unless the parent died.
            move |input_receiver, output_sender| {
                // These are `std` channels notify will send events to
                let (notify_sender, notify_receiver) = mpsc::channel();
                // These are the corresponding crossbeam channels
                let (watcher_sender, watcher_receiver) = unbounded();
                let watcher = notify::watcher(notify_sender, WATCHER_DELAY)
                    .map_err(|e| log::error!("failed to spawn notify {}", e))
                    .ok();
                let mut ctx = WatcherCtx {
                    roots,
                    watcher,
                    sender: output_sender,
                };
                let thread = thread::spawn(move || {
                    let _ = notify_receiver
                        .into_iter()
                        // forward relevant events only
                        .for_each(|event| convert_notify_event(event, &watcher_sender));
                });

                loop {
                    select! {
                        // Received request from the caller. If this channel is
                        // closed, we should shutdown everything.
                        recv(input_receiver) -> t => match t {
                            Err(RecvError) => break,
                            Ok(Task::AddRoot { root, config }) => watch_root(&mut ctx, root, Arc::clone(&config)),
                        },
                        // Watcher send us changes. If **this** channel is
                        // closed, the watcher has died, which indicates a bug
                        // -- escalate!
                        recv(watcher_receiver) -> event => match event {
                            Err(RecvError) => panic!("watcher is dead"),
                            Ok((path, change)) => WatcherCtx::handle_change(&mut ctx, path, change).unwrap(),
                        },
                    }
                }
                drop(ctx.watcher.take());
                drop(ctx);
                let res2 = thread.join();
                match &res2 {
                    Ok(()) => log::info!("... Watcher terminated with ok"),
                    Err(_) => log::error!("... Watcher terminated with err"),
                }
                res2.unwrap();
            },
        );
        Worker {
            worker,
            worker_handle,
        }
    }

    pub(crate) fn sender(&self) -> &Sender<Task> {
        &self.worker.inp
    }

    pub(crate) fn receiver(&self) -> &Receiver<TaskResult> {
        &self.worker.out
    }

    pub(crate) fn shutdown(self) -> thread::Result<()> {
        let _ = self.worker.shutdown();
        self.worker_handle.shutdown()
    }
}

fn watch_root(woker: &mut WatcherCtx, root: VfsRoot, config: Arc<RootConfig>) {
    log::debug!("loading {} ...", config.root.as_path().display());
    let files = watch_recursive(woker.watcher.as_mut(), config.root.as_path(), &*config)
        .into_iter()
        .filter_map(|path| {
            let abs_path = path.to_path(&config.root);
            let text = read_to_string(&abs_path)?;
            Some((path, text))
        })
        .collect();
    woker
        .sender
        .send(TaskResult::BulkLoadRoot { root, files })
        .unwrap();
    log::debug!("... loaded {}", config.root.as_path().display());
}

struct WatcherCtx {
    roots: Arc<Roots>,
    watcher: Option<RecommendedWatcher>,
    sender: Sender<TaskResult>,
}

fn convert_notify_event(event: DebouncedEvent, sender: &Sender<(PathBuf, ChangeKind)>) {
    match event {
        DebouncedEvent::NoticeWrite(_)
        | DebouncedEvent::NoticeRemove(_)
        | DebouncedEvent::Chmod(_) => {
            // ignore
        }
        DebouncedEvent::Rescan => {
            // TODO rescan all roots
        }
        DebouncedEvent::Create(path) => {
            sender.send((path, ChangeKind::Create)).unwrap();
        }
        DebouncedEvent::Write(path) => {
            sender.send((path, ChangeKind::Write)).unwrap();
        }
        DebouncedEvent::Remove(path) => {
            sender.send((path, ChangeKind::Remove)).unwrap();
        }
        DebouncedEvent::Rename(src, dst) => {
            sender.send((src, ChangeKind::Remove)).unwrap();
            sender.send((dst, ChangeKind::Create)).unwrap();
        }
        DebouncedEvent::Error(err, path) => {
            // TODO should we reload the file contents?
            log::warn!("watcher error \"{}\", {:?}", err, path);
        }
    }
}

impl WatcherCtx {
    fn handle_change(&mut self, path: PathBuf, kind: ChangeKind) -> Result<()> {
        let (root, rel_path) = match self.roots.find(&path) {
            None => return Ok(()),
            Some(it) => it,
        };
        let config = &self.roots[root];
        match kind {
            ChangeKind::Create => {
                let mut paths = Vec::new();
                if path.is_dir() {
                    paths.extend(watch_recursive(self.watcher.as_mut(), &path, &config));
                } else {
                    paths.push(rel_path);
                }
                paths
                    .into_iter()
                    .filter_map(|rel_path| {
                        let abs_path = rel_path.to_path(&config.root);
                        let text = read_to_string(&abs_path)?;
                        Some((rel_path, text))
                    })
                    .try_for_each(|(path, text)| {
                        self.sender
                            .send(TaskResult::AddSingleFile { root, path, text })
                    })?
            }
            ChangeKind::Write => {
                if let Some(text) = read_to_string(&path) {
                    self.sender.send(TaskResult::ChangeSingleFile {
                        root,
                        path: rel_path,
                        text,
                    })?;
                }
            }
            ChangeKind::Remove => self.sender.send(TaskResult::RemoveSingleFile {
                root,
                path: rel_path,
            })?,
        }
        Ok(())
    }
}

fn watch_recursive(
    mut watcher: Option<&mut RecommendedWatcher>,
    dir: &Path,
    config: &RootConfig,
) -> Vec<RelativePathBuf> {
    let mut files = Vec::new();
    for entry in WalkDir::new(dir)
        .into_iter()
        .filter_entry(|it| config.contains(it.path()).is_some())
        .filter_map(|it| it.map_err(|e| log::warn!("watcher error: {}", e)).ok())
    {
        if entry.file_type().is_dir() {
            if let Some(watcher) = &mut watcher {
                watch_one(watcher, entry.path());
            }
        } else {
            let path = config.contains(entry.path()).unwrap();
            files.push(path.to_owned());
        }
    }
    files
}

fn watch_one(watcher: &mut RecommendedWatcher, dir: &Path) {
    match watcher.watch(dir, RecursiveMode::NonRecursive) {
        Ok(()) => log::debug!("watching \"{}\"", dir.display()),
        Err(e) => log::warn!("could not watch \"{}\": {}", dir.display(), e),
    }
}

fn read_to_string(path: &Path) -> Option<String> {
    fs::read_to_string(&path)
        .map_err(|e| log::warn!("failed to read file {}", e))
        .ok()
}
<commit_msg>remove watcher ctx<commit_after>use std::{
    fs,
    thread,
    path::{Path, PathBuf},
    sync::{mpsc, Arc},
    time::Duration,
};
use crossbeam_channel::{Receiver, Sender, unbounded, RecvError, select};
use relative_path::RelativePathBuf;
use thread_worker::WorkerHandle;
use walkdir::WalkDir;
use notify::{DebouncedEvent, RecommendedWatcher, RecursiveMode, Watcher as _Watcher};

use crate::{RootConfig, Roots, VfsRoot};

pub(crate) enum Task {
    AddRoot {
        root: VfsRoot,
        config: Arc<RootConfig>,
    },
}

#[derive(Debug)]
pub enum TaskResult {
    BulkLoadRoot {
        root: VfsRoot,
        files: Vec<(RelativePathBuf, String)>,
    },
    AddSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
    },
    ChangeSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
        text: String,
    },
    RemoveSingleFile {
        root: VfsRoot,
        path: RelativePathBuf,
    },
}

#[derive(Debug)]
enum ChangeKind {
    Create,
    Write,
    Remove,
}

const WATCHER_DELAY: Duration = Duration::from_millis(250);

pub(crate) struct Worker {
    worker: thread_worker::Worker<Task, TaskResult>,
    worker_handle: WorkerHandle,
}

impl Worker {
    pub(crate) fn start(roots: Arc<Roots>) -> Worker {
        // This is a pretty elaborate setup of threads & channels! It is
        // explained by the following concerns:
        //    * we need to burn a thread translating from notify's mpsc to
        //      crossbeam_channel.
        //    * we want to read all files from a single thread, to gurantee that
        //      we always get fresher versions and never go back in time.
        //    * we want to tear down everything neatly during shutdown.
        let (worker, worker_handle) = thread_worker::spawn(
            "vfs",
            128,
            // This are the channels we use to communicate with outside world.
            // If `input_receiver` is closed we need to tear ourselves down.
            // `output_sender` should not be closed unless the parent died.
            move |input_receiver, output_sender| {
                // These are `std` channels notify will send events to
                let (notify_sender, notify_receiver) = mpsc::channel();
                // These are the corresponding crossbeam channels
                let (watcher_sender, watcher_receiver) = unbounded();

                let mut watcher = notify::watcher(notify_sender, WATCHER_DELAY)
                    .map_err(|e| log::error!("failed to spawn notify {}", e))
                    .ok();
                // Start a silly thread to tranform between two channels
                let thread = thread::spawn(move || {
                    notify_receiver
                        .into_iter()
                        .for_each(|event| convert_notify_event(event, &watcher_sender))
                });

                // Process requests from the called or notifications from
                // watcher until the caller says stop.
                loop {
                    select! {
                        // Received request from the caller. If this channel is
                        // closed, we should shutdown everything.
                        recv(input_receiver) -> t => match t {
                            Err(RecvError) => {
                                drop(input_receiver);
                                break
                            },
                            Ok(Task::AddRoot { root, config }) => {
                                watch_root(watcher.as_mut(), &output_sender, root, Arc::clone(&config));
                            }
                        },
                        // Watcher send us changes. If **this** channel is
                        // closed, the watcher has died, which indicates a bug
                        // -- escalate!
                        recv(watcher_receiver) -> event => match event {
                            Err(RecvError) => panic!("watcher is dead"),
                            Ok((path, change)) => {
                                handle_change(watcher.as_mut(), &output_sender, &*roots, path, change);
                            }
                        },
                    }
                }
                // Stopped the watcher
                drop(watcher.take());
                // Drain pending events: we are not inrerested in them anyways!
                watcher_receiver.into_iter().for_each(|_| ());

                let res = thread.join();
                match &res {
                    Ok(()) => log::info!("... Watcher terminated with ok"),
                    Err(_) => log::error!("... Watcher terminated with err"),
                }
                res.unwrap();
            },
        );

        Worker {
            worker,
            worker_handle,
        }
    }

    pub(crate) fn sender(&self) -> &Sender<Task> {
        &self.worker.inp
    }

    pub(crate) fn receiver(&self) -> &Receiver<TaskResult> {
        &self.worker.out
    }

    pub(crate) fn shutdown(self) -> thread::Result<()> {
        let _ = self.worker.shutdown();
        self.worker_handle.shutdown()
    }
}

fn watch_root(
    watcher: Option<&mut RecommendedWatcher>,
    sender: &Sender<TaskResult>,
    root: VfsRoot,
    config: Arc<RootConfig>,
) {
    log::debug!("loading {} ...", config.root.as_path().display());
    let files = watch_recursive(watcher, config.root.as_path(), &*config)
        .into_iter()
        .filter_map(|path| {
            let abs_path = path.to_path(&config.root);
            let text = read_to_string(&abs_path)?;
            Some((path, text))
        })
        .collect();
    sender
        .send(TaskResult::BulkLoadRoot { root, files })
        .unwrap();
    log::debug!("... loaded {}", config.root.as_path().display());
}

fn convert_notify_event(event: DebouncedEvent, sender: &Sender<(PathBuf, ChangeKind)>) {
    // forward relevant events only
    match event {
        DebouncedEvent::NoticeWrite(_)
        | DebouncedEvent::NoticeRemove(_)
        | DebouncedEvent::Chmod(_) => {
            // ignore
        }
        DebouncedEvent::Rescan => {
            // TODO rescan all roots
        }
        DebouncedEvent::Create(path) => {
            sender.send((path, ChangeKind::Create)).unwrap();
        }
        DebouncedEvent::Write(path) => {
            sender.send((path, ChangeKind::Write)).unwrap();
        }
        DebouncedEvent::Remove(path) => {
            sender.send((path, ChangeKind::Remove)).unwrap();
        }
        DebouncedEvent::Rename(src, dst) => {
            sender.send((src, ChangeKind::Remove)).unwrap();
            sender.send((dst, ChangeKind::Create)).unwrap();
        }
        DebouncedEvent::Error(err, path) => {
            // TODO should we reload the file contents?
            log::warn!("watcher error \"{}\", {:?}", err, path);
        }
    }
}

fn handle_change(
    watcher: Option<&mut RecommendedWatcher>,
    sender: &Sender<TaskResult>,
    roots: &Roots,
    path: PathBuf,
    kind: ChangeKind,
) {
    let (root, rel_path) = match roots.find(&path) {
        None => return,
        Some(it) => it,
    };
    let config = &roots[root];
    match kind {
        ChangeKind::Create => {
            let mut paths = Vec::new();
            if path.is_dir() {
                paths.extend(watch_recursive(watcher, &path, &config));
            } else {
                paths.push(rel_path);
            }
            paths
                .into_iter()
                .filter_map(|rel_path| {
                    let abs_path = rel_path.to_path(&config.root);
                    let text = read_to_string(&abs_path)?;
                    Some((rel_path, text))
                })
                .try_for_each(|(path, text)| {
                    sender.send(TaskResult::AddSingleFile { root, path, text })
                })
                .unwrap()
        }
        ChangeKind::Write => {
            if let Some(text) = read_to_string(&path) {
                sender
                    .send(TaskResult::ChangeSingleFile {
                        root,
                        path: rel_path,
                        text,
                    })
                    .unwrap();
            }
        }
        ChangeKind::Remove => sender
            .send(TaskResult::RemoveSingleFile {
                root,
                path: rel_path,
            })
            .unwrap(),
    }
}

fn watch_recursive(
    mut watcher: Option<&mut RecommendedWatcher>,
    dir: &Path,
    config: &RootConfig,
) -> Vec<RelativePathBuf> {
    let mut files = Vec::new();
    for entry in WalkDir::new(dir)
        .into_iter()
        .filter_entry(|it| config.contains(it.path()).is_some())
        .filter_map(|it| it.map_err(|e| log::warn!("watcher error: {}", e)).ok())
    {
        if entry.file_type().is_dir() {
            if let Some(watcher) = &mut watcher {
                watch_one(watcher, entry.path());
            }
        } else {
            let path = config.contains(entry.path()).unwrap();
            files.push(path.to_owned());
        }
    }
    files
}

fn watch_one(watcher: &mut RecommendedWatcher, dir: &Path) {
    match watcher.watch(dir, RecursiveMode::NonRecursive) {
        Ok(()) => log::debug!("watching \"{}\"", dir.display()),
        Err(e) => log::warn!("could not watch \"{}\": {}", dir.display(), e),
    }
}

fn read_to_string(path: &Path) -> Option<String> {
    fs::read_to_string(&path)
        .map_err(|e| log::warn!("failed to read file {}", e))
        .ok()
}
<|endoftext|>