<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 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 IPv4 addresses
    fn is_ipv4_empty(&self) -> bool {
        self.ipv4.iter().next().is_none()
    }

    /// Check if there are no rules for IPv6 addresses
    fn is_ipv6_empty(&self) -> bool {
        self.ipv6.iter().next().is_none()
    }
}

/// 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_host_in_proxy_list(&self, host: &str) -> Option<bool> {
        // Addresses in proxy_list will be proxied
        if self.white_list.check_host_matched(host) {
            return Some(true);
        }
        // Addresses in bypass_list will be bypassed
        if self.black_list.check_host_matched(host) {
            return Some(false);
        }
        None
    }

    /// If there are no IPv4 rules
    pub fn is_ipv4_empty(&self) -> bool {
        match self.mode {
            Mode::BlackList => self.black_list.is_ipv4_empty(),
            Mode::WhiteList => self.white_list.is_ipv4_empty(),
        }
    }

    /// If there are no IPv4 rules
    pub fn is_ipv6_empty(&self) -> bool {
        match self.mode {
            Mode::BlackList => self.black_list.is_ipv6_empty(),
            Mode::WhiteList => self.white_list.is_ipv6_empty(),
        }
    }

    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),
        }
    }

    pub 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)
    ///
    /// This function may perform a DNS resolution
    pub async fn check_target_bypassed(&self, context: &Context, addr: &Address) -> bool {
        match *addr {
            Address::SocketAddress(ref addr) => self.check_ip_in_proxy_list(&addr.ip()),
            // Resolve hostname and check the list
            Address::DomainNameAddress(ref host, port) => {
                if let Some(value) = self.check_host_in_proxy_list(host) {
                    return value;
                }
                if let Ok(vaddr) = context.dns_resolve(host, port).await {
                    for addr in vaddr {
                        if !self.check_ip_in_proxy_list(&addr.ip()) {
                            return true;
                        }
                    }
                }
                false
            },
        }
    }

    /// 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>Bypass dns_resolve if IP rules are empty<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 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 IPv4 addresses
    fn is_ipv4_empty(&self) -> bool {
        self.ipv4.iter().next().is_none()
    }

    /// Check if there are no rules for IPv6 addresses
    fn is_ipv6_empty(&self) -> bool {
        self.ipv6.iter().next().is_none()
    }
}

/// 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_host_in_proxy_list(&self, host: &str) -> Option<bool> {
        // Addresses in proxy_list will be proxied
        if self.white_list.check_host_matched(host) {
            return Some(true);
        }
        // Addresses in bypass_list will be bypassed
        if self.black_list.check_host_matched(host) {
            return Some(false);
        }
        None
    }

    /// If there are no IPv4 rules
    pub fn is_ipv4_empty(&self) -> bool {
        match self.mode {
            Mode::BlackList => self.black_list.is_ipv4_empty(),
            Mode::WhiteList => self.white_list.is_ipv4_empty(),
        }
    }

    /// If there are no IPv4 rules
    pub fn is_ipv6_empty(&self) -> bool {
        match self.mode {
            Mode::BlackList => self.black_list.is_ipv6_empty(),
            Mode::WhiteList => self.white_list.is_ipv6_empty(),
        }
    }

    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),
        }
    }

    pub 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)
    ///
    /// This function may perform a DNS resolution
    pub async fn check_target_bypassed(&self, context: &Context, addr: &Address) -> bool {
        match *addr {
            Address::SocketAddress(ref addr) => self.check_ip_in_proxy_list(&addr.ip()),
            // Resolve hostname and check the list
            Address::DomainNameAddress(ref host, port) => {
                if let Some(value) = self.check_host_in_proxy_list(host) {
                    return value;
                }
                if self.is_ipv4_empty() && self.is_ipv6_empty() {
                    return !self.is_default_in_proxy_list();
                }
                if let Ok(vaddr) = context.dns_resolve(host, port).await {
                    for addr in vaddr {
                        if !self.check_ip_in_proxy_list(&addr.ip()) {
                            return true;
                        }
                    }
                }
                false
            },
        }
    }

    /// 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>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use super::block::SBBlock;
use super::compileunit::SBCompileUnit;
use super::function::SBFunction;
use super::lineentry::SBLineEntry;
use super::module::SBModule;
use super::stream::SBStream;
use super::symbol::SBSymbol;
use super::symbolcontext::SBSymbolContext;
use super::target::SBTarget;
use std::fmt;
use sys;

/// A section + offset based address class.
///
/// The `SBAddress` class allows addresses to be relative to a section
/// that can move during runtime due to images (executables, shared
/// libraries, bundles, frameworks) being loaded at different
/// addresses than the addresses found in the object file that
/// represents them on disk. There are currently two types of addresses
/// for a section:
///
/// * file addresses
/// * load addresses
///
/// File addresses represents the virtual addresses that are in the 'on
/// disk' object files. These virtual addresses are converted to be
/// relative to unique sections scoped to the object file so that
/// when/if the addresses slide when the images are loaded/unloaded
/// in memory, we can easily track these changes without having to
/// update every object (compile unit ranges, line tables, function
/// address ranges, lexical block and inlined subroutine address
/// ranges, global and static variables) each time an image is loaded or
/// unloaded.
///
/// Load addresses represents the virtual addresses where each section
/// ends up getting loaded at runtime. Before executing a program, it
/// is common for all of the load addresses to be unresolved. When a
/// `DynamicLoader` plug-in receives notification that shared libraries
/// have been loaded/unloaded, the load addresses of the main executable
/// and any images (shared libraries) will be  resolved/unresolved. When
/// this happens, breakpoints that are in one of these sections can be
/// set/cleared.
pub struct SBAddress {
    /// The underlying raw `SBAddressRef`.
    pub raw: sys::SBAddressRef,
}

impl SBAddress {
    /// Construct a new `SBAddress`.
    pub fn wrap(raw: sys::SBAddressRef) -> SBAddress {
        SBAddress { raw }
    }

    /// Construct a new `Some(SBAddress)` or `None`.
    pub fn maybe_wrap(raw: sys::SBAddressRef) -> Option<SBAddress> {
        if unsafe { sys::SBAddressIsValid(raw) != 0 } {
            Some(SBAddress { raw })
        } else {
            None
        }
    }

    /// Check whether or not this is a valid `SBAddress` value.
    pub fn is_valid(&self) -> bool {
        unsafe { sys::SBAddressIsValid(self.raw) != 0 }
    }

    /// The address that represents the address as it is found in the
    /// object file that defines it.
    pub fn file_address(&self) -> u64 {
        unsafe { sys::SBAddressGetFileAddress(self.raw) }
    }

    /// The address as it has been loaded into memory by a target.
    pub fn load_address(&self, target: &SBTarget) -> u64 {
        unsafe { sys::SBAddressGetLoadAddress(self.raw, target.raw) }
    }

    /// Get the `SBSymbolContext` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// * `resolve_scope`: Flags that specify what type of symbol context
    ///   is needed by the caller. These flags have constants starting
    ///   with `SYMBOL_CONTEXT_ITEM_`.
    pub fn symbol_context(&self, resolve_scope: u32) -> SBSymbolContext {
        SBSymbolContext::wrap(unsafe { sys::SBAddressGetSymbolContext(self.raw, resolve_scope) })
    }

    /// Get the `SBModule` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn module(&self) -> Option<SBModule> {
        SBModule::maybe_wrap(unsafe { sys::SBAddressGetModule(self.raw) })
    }

    /// Get the `SBCompileUnit` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn compile_unit(&self) -> Option<SBCompileUnit> {
        SBCompileUnit::maybe_wrap(unsafe { sys::SBAddressGetCompileUnit(self.raw) })
    }

    /// Get the `SBFunction` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn function(&self) -> Option<SBFunction> {
        SBFunction::maybe_wrap(unsafe { sys::SBAddressGetFunction(self.raw) })
    }

    /// Get the `SBBlock` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn block(&self) -> Option<SBBlock> {
        SBBlock::maybe_wrap(unsafe { sys::SBAddressGetBlock(self.raw) })
    }

    /// Get the `SBSymbol` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn symbol(&self) -> Option<SBSymbol> {
        SBSymbol::maybe_wrap(unsafe { sys::SBAddressGetSymbol(self.raw) })
    }

    /// Get the `SBLineEntry` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn line_entry(&self) -> Option<SBLineEntry> {
        SBLineEntry::maybe_wrap(unsafe { sys::SBAddressGetLineEntry(self.raw) })
    }
}

impl Clone for SBAddress {
    fn clone(&self) -> SBAddress {
        SBAddress {
            raw: unsafe { sys::CloneSBAddress(self.raw) },
        }
    }
}

impl fmt::Debug for SBAddress {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let stream = SBStream::new();
        unsafe { sys::SBAddressGetDescription(self.raw, stream.raw) };
        write!(fmt, "SBAddress {{ {} }}", stream.data())
    }
}

impl Drop for SBAddress {
    fn drop(&mut self) {
        unsafe { sys::DisposeSBAddress(self.raw) };
    }
}

unsafe impl Send for SBAddress {}
unsafe impl Sync for SBAddress {}

#[cfg(feature = "graphql")]
graphql_object!(SBAddress: super::debugger::SBDebugger | &self | {
    field is_valid() -> bool {
        self.is_valid()
    }

    // TODO(bm) This should be u64
    field file_address() -> i32 {
        self.file_address() as i32
    }

    field module() -> Option<SBModule> {
        self.module()
    }

    field compile_unit() -> Option<SBCompileUnit> {
        self.compile_unit()
    }

    field function() -> Option<SBFunction> {
        self.function()
    }

    field block() -> Option<SBBlock> {
        self.block()
    }

    field symbol() -> Option<SBSymbol> {
        self.symbol()
    }

    field line_entry() -> Option<SBLineEntry> {
        self.line_entry()
    }
});
<commit_msg>Implement Eq and PartialEq for SBAddress<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use super::block::SBBlock;
use super::compileunit::SBCompileUnit;
use super::function::SBFunction;
use super::lineentry::SBLineEntry;
use super::module::SBModule;
use super::stream::SBStream;
use super::symbol::SBSymbol;
use super::symbolcontext::SBSymbolContext;
use super::target::SBTarget;
use std::fmt;
use sys;

/// A section + offset based address class.
///
/// The `SBAddress` class allows addresses to be relative to a section
/// that can move during runtime due to images (executables, shared
/// libraries, bundles, frameworks) being loaded at different
/// addresses than the addresses found in the object file that
/// represents them on disk. There are currently two types of addresses
/// for a section:
///
/// * file addresses
/// * load addresses
///
/// File addresses represents the virtual addresses that are in the 'on
/// disk' object files. These virtual addresses are converted to be
/// relative to unique sections scoped to the object file so that
/// when/if the addresses slide when the images are loaded/unloaded
/// in memory, we can easily track these changes without having to
/// update every object (compile unit ranges, line tables, function
/// address ranges, lexical block and inlined subroutine address
/// ranges, global and static variables) each time an image is loaded or
/// unloaded.
///
/// Load addresses represents the virtual addresses where each section
/// ends up getting loaded at runtime. Before executing a program, it
/// is common for all of the load addresses to be unresolved. When a
/// `DynamicLoader` plug-in receives notification that shared libraries
/// have been loaded/unloaded, the load addresses of the main executable
/// and any images (shared libraries) will be  resolved/unresolved. When
/// this happens, breakpoints that are in one of these sections can be
/// set/cleared.
pub struct SBAddress {
    /// The underlying raw `SBAddressRef`.
    pub raw: sys::SBAddressRef,
}

impl SBAddress {
    /// Construct a new `SBAddress`.
    pub fn wrap(raw: sys::SBAddressRef) -> SBAddress {
        SBAddress { raw }
    }

    /// Construct a new `Some(SBAddress)` or `None`.
    pub fn maybe_wrap(raw: sys::SBAddressRef) -> Option<SBAddress> {
        if unsafe { sys::SBAddressIsValid(raw) != 0 } {
            Some(SBAddress { raw })
        } else {
            None
        }
    }

    /// Check whether or not this is a valid `SBAddress` value.
    pub fn is_valid(&self) -> bool {
        unsafe { sys::SBAddressIsValid(self.raw) != 0 }
    }

    /// The address that represents the address as it is found in the
    /// object file that defines it.
    pub fn file_address(&self) -> u64 {
        unsafe { sys::SBAddressGetFileAddress(self.raw) }
    }

    /// The address as it has been loaded into memory by a target.
    pub fn load_address(&self, target: &SBTarget) -> u64 {
        unsafe { sys::SBAddressGetLoadAddress(self.raw, target.raw) }
    }

    /// Get the `SBSymbolContext` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// * `resolve_scope`: Flags that specify what type of symbol context
    ///   is needed by the caller. These flags have constants starting
    ///   with `SYMBOL_CONTEXT_ITEM_`.
    pub fn symbol_context(&self, resolve_scope: u32) -> SBSymbolContext {
        SBSymbolContext::wrap(unsafe { sys::SBAddressGetSymbolContext(self.raw, resolve_scope) })
    }

    /// Get the `SBModule` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn module(&self) -> Option<SBModule> {
        SBModule::maybe_wrap(unsafe { sys::SBAddressGetModule(self.raw) })
    }

    /// Get the `SBCompileUnit` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn compile_unit(&self) -> Option<SBCompileUnit> {
        SBCompileUnit::maybe_wrap(unsafe { sys::SBAddressGetCompileUnit(self.raw) })
    }

    /// Get the `SBFunction` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn function(&self) -> Option<SBFunction> {
        SBFunction::maybe_wrap(unsafe { sys::SBAddressGetFunction(self.raw) })
    }

    /// Get the `SBBlock` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn block(&self) -> Option<SBBlock> {
        SBBlock::maybe_wrap(unsafe { sys::SBAddressGetBlock(self.raw) })
    }

    /// Get the `SBSymbol` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn symbol(&self) -> Option<SBSymbol> {
        SBSymbol::maybe_wrap(unsafe { sys::SBAddressGetSymbol(self.raw) })
    }

    /// Get the `SBLineEntry` for a given address.
    ///
    /// An address might refer to code or data from an existing
    /// module, or it might refer to something on the stack or heap.
    /// This will only return valid values if the address has been
    /// resolved to a code or data address using
    /// `SBAddress::set_load_address` or `SBTarget::resolve_load_address`.
    ///
    /// This grabs an individual object for a given address and
    /// is less efficient if you want more than one symbol related objects.
    /// Use one of the following when you want multiple debug symbol related
    /// objects for an address:
    ///
    /// * `SBAddress::symbol_context`
    /// * `SBTarget::resolve_symbol_context_for_address`
    ///
    /// One or more bits from the `SymbolContextItem` enumerations can be logically
    /// OR'ed together to more efficiently retrieve multiple symbol objects.
    pub fn line_entry(&self) -> Option<SBLineEntry> {
        SBLineEntry::maybe_wrap(unsafe { sys::SBAddressGetLineEntry(self.raw) })
    }
}

impl Clone for SBAddress {
    fn clone(&self) -> SBAddress {
        SBAddress {
            raw: unsafe { sys::CloneSBAddress(self.raw) },
        }
    }
}

impl fmt::Debug for SBAddress {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let stream = SBStream::new();
        unsafe { sys::SBAddressGetDescription(self.raw, stream.raw) };
        write!(fmt, "SBAddress {{ {} }}", stream.data())
    }
}

impl Drop for SBAddress {
    fn drop(&mut self) {
        unsafe { sys::DisposeSBAddress(self.raw) };
    }
}

unsafe impl Send for SBAddress {}
unsafe impl Sync for SBAddress {}

impl PartialEq for SBAddress {
    fn eq(&self, other: &Self) -> bool {
        unsafe { sys::SBAddressIsEqual(self.raw, other.raw) }
    }
}

impl Eq for SBAddress {}

#[cfg(feature = "graphql")]
graphql_object!(SBAddress: super::debugger::SBDebugger | &self | {
    field is_valid() -> bool {
        self.is_valid()
    }

    // TODO(bm) This should be u64
    field file_address() -> i32 {
        self.file_address() as i32
    }

    field module() -> Option<SBModule> {
        self.module()
    }

    field compile_unit() -> Option<SBCompileUnit> {
        self.compile_unit()
    }

    field function() -> Option<SBFunction> {
        self.function()
    }

    field block() -> Option<SBBlock> {
        self.block()
    }

    field symbol() -> Option<SBSymbol> {
        self.symbol()
    }

    field line_entry() -> Option<SBLineEntry> {
        self.line_entry()
    }
});

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

    #[test]
    fn test_equal() {
        let sect = unsafe { sys::CreateSBSection() };
        let a = SBAddress::maybe_wrap(unsafe { sys::CreateSBAddress2(sect, 42) }).unwrap();
        let b = SBAddress::maybe_wrap(unsafe { sys::CreateSBAddress2(sect, 42) }).unwrap();
        assert!(a == b);
    }

    #[test]
    fn test_not_equal() {
        let sect = unsafe { sys::CreateSBSection() };
        let a = SBAddress::maybe_wrap(unsafe { sys::CreateSBAddress2(sect, 42) }).unwrap();
        let b = SBAddress::maybe_wrap(unsafe { sys::CreateSBAddress2(sect, 111) }).unwrap();
        assert!(a != b);
    }
}
<|endoftext|>


<commit_before>use rustc_serialize::{Encodable, Decodable};
use super::master::{self, Master};
use super::slave::Slave;
use super::error::{ErrorKind, Result};
use super::error::master::Result as MasterResult;
use super::value::Topic;
use super::naming::{self, Resolver};
use super::resolve;
use tcpros::{Client, Message, PublisherStream, ServicePair, ServiceResult};
use rosxmlrpc::serde::XmlRpcValue;
use ::std::time::Duration;

pub struct Ros {
    master: Master,
    slave: Slave,
    hostname: String,
    resolver: Resolver,
    name: String,
}

impl Ros {
    pub fn new(name: &str) -> Result<Ros> {
        let namespace = resolve::namespace();
        let master_uri = resolve::master();
        let hostname = resolve::hostname();
        let name = resolve::name(name);
        let mut ros = Ros::new_raw(&master_uri, &hostname, &namespace, &name)?;
        for (src, dest) in resolve::mappings() {
            ros.map(&src, &dest)?;
        }
        Ok(ros)
    }

    fn new_raw(master_uri: &str, hostname: &str, namespace: &str, name: &str) -> Result<Ros> {
        let namespace = namespace.trim_right_matches("/");

        if name.contains("/") {
            bail!(ErrorKind::Naming(naming::error::ErrorKind::IllegalCharacter(name.into())));
        }

        let name = format!("{}/{}", namespace, name);
        let resolver = Resolver::new(&name)?;

        let slave = Slave::new(&master_uri, &hostname, 0, &name)?;
        let master = Master::new(&master_uri, &name, &slave.uri());
        Ok(Ros {
               master: master,
               slave: slave,
               hostname: String::from(hostname),
               resolver: resolver,
               name: name,
           })
    }

    fn map(&mut self, source: &str, destination: &str) -> Result<()> {
        self.resolver.map(source, destination).map_err(|v| v.into())
    }

    pub fn uri(&self) -> &str {
        return self.slave.uri();
    }

    pub fn param<'a, 'b>(&'a self, name: &'b str) -> Option<Parameter<'a>> {
        self.resolver
            .translate(name)
            .ok()
            .map(|v| {
                     Parameter {
                         master: &self.master,
                         name: v,
                     }
                 })
    }

    pub fn parameters(&self) -> MasterResult<Vec<String>> {
        self.master.get_param_names()
    }

    pub fn state(&self) -> MasterResult<master::SystemState> {
        self.master.get_system_state()
    }

    pub fn topics(&self) -> MasterResult<Vec<Topic>> {
        self.master.get_topic_types()
    }

    pub fn client<T: ServicePair>(&self, service: &str) -> Result<Client<T>> {
        let name = self.resolver.translate(service)?;
        let uri = self.master.lookup_service(&name)?;
        Ok(Client::new(&self.name, &uri, &name))
    }

    pub fn wait_for_service(&self, service: &str, timeout: Option<Duration>) -> Result<()> {
        use super::error::master::ErrorKind::Api;
        use super::error::api::ErrorKind::BadData;
        use ::std::thread::sleep;

        let name = self.resolver.translate(service)?;
        let now = ::std::time::Instant::now();
        loop {
            match self.master.lookup_service(&name) {
                Err(e) => {
                    match e.kind() {
                        &Api(BadData(ref m)) if m == "no provider" => {
			                if let Some(ref timeout) = timeout {
			                    if &now.elapsed() > timeout {
				                    return Err(ErrorKind::TimeoutError.into());
				                }
			                }
                            sleep(Duration::from_millis(100));
                            continue;
                        },
                        _ => {},
                    }
                    return Err(e.into());
                }
                Ok(_) => {break}
            }
        }
        Ok(())
    }

    pub fn service<T, F>(&mut self, service: &str, handler: F) -> Result<()>
        where T: ServicePair,
              F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static
    {
        let name = self.resolver.translate(service)?;
        let api = self.slave
            .add_service::<T, F>(&self.hostname, &name, handler)?;

        if let Err(err) = self.master.register_service(&name, &api) {
            self.slave.remove_service(&name);
            self.master.unregister_service(&name, &api)?;
            Err(err.into())
        } else {
            Ok(())
        }

    }

    pub fn subscribe<T, F>(&mut self, topic: &str, callback: F) -> Result<()>
        where T: Message,
              F: Fn(T) -> () + Send + 'static
    {
        let name = self.resolver.translate(topic)?;
        self.slave.add_subscription::<T, F>(&name, callback)?;

        match self.master.register_subscriber(&name, &T::msg_type()) {
            Ok(publishers) => {
                if let Err(err) =
                    self.slave
                        .add_publishers_to_subscription(&name, publishers.into_iter()) {
                    error!("Failed to subscribe to all publishers of topic '{}': {}",
                           name,
                           err);
                }
                Ok(())
            }
            Err(err) => {
                self.slave.remove_subscription(&name);
                self.master.unregister_subscriber(&name)?;
                Err(err.into())
            }
        }
    }

    pub fn publish<T>(&mut self, topic: &str) -> Result<PublisherStream<T>>
        where T: Message
    {
        let name = self.resolver.translate(topic)?;
        let stream = self.slave.add_publication::<T>(&self.hostname, &name)?;
        match self.master.register_publisher(&name, &T::msg_type()) {
            Ok(_) => Ok(stream),
            Err(error) => {
                error!("Failed to register publisher for topic '{}': {}",
                       name,
                       error);
                self.slave.remove_publication(&name);
                self.master.unregister_publisher(&name)?;
                Err(error.into())
            }
        }
    }
}

pub struct Parameter<'a> {
    master: &'a Master,
    name: String,
}

impl<'a> Parameter<'a> {
    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn get<T: Decodable>(&self) -> MasterResult<T> {
        self.master.get_param::<T>(&self.name)
    }

    pub fn get_raw(&self) -> MasterResult<XmlRpcValue> {
        self.master.get_param_any(&self.name)
    }

    pub fn set<T: Encodable>(&self, value: &T) -> MasterResult<()> {
        self.master.set_param::<T>(&self.name, value).and(Ok(()))
    }

    pub fn delete(&self) -> MasterResult<()> {
        self.master.delete_param(&self.name).and(Ok(()))
    }

    pub fn exists(&self) -> MasterResult<bool> {
        self.master.has_param(&self.name)
    }

    pub fn search(&self) -> MasterResult<String> {
        self.master.search_param(&self.name)
    }
}
<commit_msg>formatting fixes<commit_after>use rustc_serialize::{Encodable, Decodable};
use super::master::{self, Master};
use super::slave::Slave;
use super::error::{ErrorKind, Result};
use super::error::master::Result as MasterResult;
use super::value::Topic;
use super::naming::{self, Resolver};
use super::resolve;
use tcpros::{Client, Message, PublisherStream, ServicePair, ServiceResult};
use rosxmlrpc::serde::XmlRpcValue;
use std::time::Duration;

pub struct Ros {
    master: Master,
    slave: Slave,
    hostname: String,
    resolver: Resolver,
    name: String,
}

impl Ros {
    pub fn new(name: &str) -> Result<Ros> {
        let namespace = resolve::namespace();
        let master_uri = resolve::master();
        let hostname = resolve::hostname();
        let name = resolve::name(name);
        let mut ros = Ros::new_raw(&master_uri, &hostname, &namespace, &name)?;
        for (src, dest) in resolve::mappings() {
            ros.map(&src, &dest)?;
        }
        Ok(ros)
    }

    fn new_raw(master_uri: &str, hostname: &str, namespace: &str, name: &str) -> Result<Ros> {
        let namespace = namespace.trim_right_matches("/");

        if name.contains("/") {
            bail!(ErrorKind::Naming(naming::error::ErrorKind::IllegalCharacter(name.into())));
        }

        let name = format!("{}/{}", namespace, name);
        let resolver = Resolver::new(&name)?;

        let slave = Slave::new(&master_uri, &hostname, 0, &name)?;
        let master = Master::new(&master_uri, &name, &slave.uri());
        Ok(Ros {
               master: master,
               slave: slave,
               hostname: String::from(hostname),
               resolver: resolver,
               name: name,
           })
    }

    fn map(&mut self, source: &str, destination: &str) -> Result<()> {
        self.resolver.map(source, destination).map_err(|v| v.into())
    }

    pub fn uri(&self) -> &str {
        return self.slave.uri();
    }

    pub fn param<'a, 'b>(&'a self, name: &'b str) -> Option<Parameter<'a>> {
        self.resolver
            .translate(name)
            .ok()
            .map(|v| {
                     Parameter {
                         master: &self.master,
                         name: v,
                     }
                 })
    }

    pub fn parameters(&self) -> MasterResult<Vec<String>> {
        self.master.get_param_names()
    }

    pub fn state(&self) -> MasterResult<master::SystemState> {
        self.master.get_system_state()
    }

    pub fn topics(&self) -> MasterResult<Vec<Topic>> {
        self.master.get_topic_types()
    }

    pub fn client<T: ServicePair>(&self, service: &str) -> Result<Client<T>> {
        let name = self.resolver.translate(service)?;
        let uri = self.master.lookup_service(&name)?;
        Ok(Client::new(&self.name, &uri, &name))
    }

    pub fn wait_for_service(&self, service: &str, timeout: Option<Duration>) -> Result<()> {
        use super::error::master::ErrorKind::Api;
        use super::error::api::ErrorKind::BadData;
        use std::thread::sleep;

        let name = self.resolver.translate(service)?;
        let now = ::std::time::Instant::now();
        loop {
            match self.master.lookup_service(&name) {
                Err(e) => {
                    match e.kind() {
                        &Api(BadData(ref m)) if m == "no provider" => {
                            if let Some(ref timeout) = timeout {
                                if &now.elapsed() > timeout {
                                    return Err(ErrorKind::TimeoutError.into());
                                }
                            }
                            sleep(Duration::from_millis(100));
                            continue;
                        }
                        _ => {}
                    }
                    return Err(e.into());
                }
                Ok(_) => break,
            }
        }
        Ok(())
    }

    pub fn service<T, F>(&mut self, service: &str, handler: F) -> Result<()>
        where T: ServicePair,
              F: Fn(T::Request) -> ServiceResult<T::Response> + Send + Sync + 'static
    {
        let name = self.resolver.translate(service)?;
        let api = self.slave
            .add_service::<T, F>(&self.hostname, &name, handler)?;

        if let Err(err) = self.master.register_service(&name, &api) {
            self.slave.remove_service(&name);
            self.master.unregister_service(&name, &api)?;
            Err(err.into())
        } else {
            Ok(())
        }

    }

    pub fn subscribe<T, F>(&mut self, topic: &str, callback: F) -> Result<()>
        where T: Message,
              F: Fn(T) -> () + Send + 'static
    {
        let name = self.resolver.translate(topic)?;
        self.slave.add_subscription::<T, F>(&name, callback)?;

        match self.master.register_subscriber(&name, &T::msg_type()) {
            Ok(publishers) => {
                if let Err(err) =
                    self.slave
                        .add_publishers_to_subscription(&name, publishers.into_iter()) {
                    error!("Failed to subscribe to all publishers of topic '{}': {}",
                           name,
                           err);
                }
                Ok(())
            }
            Err(err) => {
                self.slave.remove_subscription(&name);
                self.master.unregister_subscriber(&name)?;
                Err(err.into())
            }
        }
    }

    pub fn publish<T>(&mut self, topic: &str) -> Result<PublisherStream<T>>
        where T: Message
    {
        let name = self.resolver.translate(topic)?;
        let stream = self.slave.add_publication::<T>(&self.hostname, &name)?;
        match self.master.register_publisher(&name, &T::msg_type()) {
            Ok(_) => Ok(stream),
            Err(error) => {
                error!("Failed to register publisher for topic '{}': {}",
                       name,
                       error);
                self.slave.remove_publication(&name);
                self.master.unregister_publisher(&name)?;
                Err(error.into())
            }
        }
    }
}

pub struct Parameter<'a> {
    master: &'a Master,
    name: String,
}

impl<'a> Parameter<'a> {
    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn get<T: Decodable>(&self) -> MasterResult<T> {
        self.master.get_param::<T>(&self.name)
    }

    pub fn get_raw(&self) -> MasterResult<XmlRpcValue> {
        self.master.get_param_any(&self.name)
    }

    pub fn set<T: Encodable>(&self, value: &T) -> MasterResult<()> {
        self.master.set_param::<T>(&self.name, value).and(Ok(()))
    }

    pub fn delete(&self) -> MasterResult<()> {
        self.master.delete_param(&self.name).and(Ok(()))
    }

    pub fn exists(&self) -> MasterResult<bool> {
        self.master.has_param(&self.name)
    }

    pub fn search(&self) -> MasterResult<String> {
        self.master.search_param(&self.name)
    }
}
<|endoftext|>


<commit_before>#![feature(link_args)]
#![feature(struct_variant)]

extern crate libc;
use libc::types::os::arch::c95::{c_int, c_uint};

pub struct RawEvent {
    etype: u8,
    emod: u8,
    key: u16,
    ch: u32,
    w: i32,
    h: i32,
}

pub enum Event {
    KeyEvent(u8, u16, u32),
    ResizeEvent(i32, i32),
    NoEvent
}

pub enum Color {
    Black,
    Red,
    Green,
    Yellow,
    Blue,
    Magenta,
    Cyan,
    White
}

pub enum Style {
    Normal,
    Bold,
    Underline,
    BoldUnderline
}

mod termbox {
    use libc::types::os::arch::c95::{c_int, c_uint};

    #[link(name="termbox")]
    extern {
        pub fn tb_init() -> c_int;
        pub fn tb_shutdown();

        pub fn tb_width() -> c_uint;
        pub fn tb_height() -> c_uint;

        pub fn tb_clear();
        pub fn tb_present();

        pub fn tb_set_cursor(cx: c_int, cy: c_int);
        pub fn tb_change_cell(x: c_uint, y: c_uint, ch: u32, fg: u16, bg: u16);

        //pub fn tb_select_input_mode(mode: c_int) -> c_int;
        //pub fn tb_set_clear_attributes(fg: u16, bg: u16);

        pub fn tb_peek_event(ev: *const ::RawEvent, timeout: c_uint) -> c_int;
        pub fn tb_poll_event(ev: *const ::RawEvent) -> c_int;
    }
}

fn nil_raw_event() -> RawEvent {
    RawEvent{etype: 0, emod: 0, key: 0, ch: 0, w: 0, h: 0}
}

fn unpack_event(ev_type: c_int, ev: &RawEvent) -> Event {
    match ev_type {
        0 => NoEvent,
        1 => {
            return KeyEvent(ev.emod, ev.key, ev.ch);
        },
        2 => {
            return ResizeEvent(ev.w, ev.h);
        },
        _ => { fail!("Unknown event"); }
    }
}

pub fn convert_color(c: Color) -> u16 {
    match c {
        Black   => 0x00,
        Red     => 0x01,
        Green   => 0x02,
        Yellow  => 0x03,
        Blue    => 0x04,
        Magenta => 0x05,
        Cyan    => 0x06,
        White   => 0x07,
    }
}

pub fn convert_style(sty: Style) -> u16 {
    match sty {
        Normal         => 0x00,
        Bold           => 0x10,
        Underline      => 0x20,
        BoldUnderline => 0x30,
    }
}

pub fn init() -> int {
    unsafe {
        termbox::tb_init() as int
    }
}

pub fn shutdown() {
    unsafe { termbox::tb_shutdown() }
}

pub fn width() -> uint {
    unsafe { termbox::tb_width() as uint }
}

pub fn height() -> uint {
    unsafe { termbox::tb_height() as uint }
}

pub fn clear() {
    unsafe { termbox::tb_clear() }
}

pub fn present() {
    unsafe { termbox::tb_present() }
}

pub fn set_cursor(x: int, y: int) {
    unsafe { termbox::tb_set_cursor(x as c_int, y as c_int) }
}

pub fn change_cell(x: uint, y: uint, ch: u32, fg: u16, bg: u16) {
    unsafe {
        termbox::tb_change_cell(x as c_uint, y as c_uint, ch, fg, bg)
    }
}

pub fn print(x: uint, y: uint, sty: Style, fg: Color, bg: Color, s: String) {
    let fg: u16 = convert_color(fg) | convert_style(sty);
    let bg: u16 = convert_color(bg);
    for (i, ch) in s.as_slice().chars().enumerate() {
        change_cell(x+i, y, ch as u32, fg, bg);
    }
}

pub fn print_char(x: uint, y: uint, sty: Style, fg: Color, bg: Color, ch: char) {
    let fg: u16 = convert_color(fg) | convert_style(sty);
    let bg: u16 = convert_color(bg);
    change_cell(x, y, ch as u32, fg, bg);
}

pub fn poll_event() -> Event {
    unsafe {
        let ev = nil_raw_event();
        let rc = termbox::tb_poll_event(&ev as *const RawEvent);
        return unpack_event(rc, &ev);
    }
}

pub fn peek_event(timeout: uint) -> Event {
    unsafe {
        let ev = nil_raw_event();
        let rc = termbox::tb_peek_event(&ev as *const RawEvent, timeout as c_uint);
        return unpack_event(rc, &ev);
    }
}
<commit_msg>Change fail to panic<commit_after>#![feature(link_args)]
#![feature(struct_variant)]

extern crate libc;
use libc::types::os::arch::c95::{c_int, c_uint};

pub struct RawEvent {
    etype: u8,
    emod: u8,
    key: u16,
    ch: u32,
    w: i32,
    h: i32,
}

pub enum Event {
    KeyEvent(u8, u16, u32),
    ResizeEvent(i32, i32),
    NoEvent
}

pub enum Color {
    Black,
    Red,
    Green,
    Yellow,
    Blue,
    Magenta,
    Cyan,
    White
}

pub enum Style {
    Normal,
    Bold,
    Underline,
    BoldUnderline
}

mod termbox {
    use libc::types::os::arch::c95::{c_int, c_uint};

    #[link(name="termbox")]
    extern {
        pub fn tb_init() -> c_int;
        pub fn tb_shutdown();

        pub fn tb_width() -> c_uint;
        pub fn tb_height() -> c_uint;

        pub fn tb_clear();
        pub fn tb_present();

        pub fn tb_set_cursor(cx: c_int, cy: c_int);
        pub fn tb_change_cell(x: c_uint, y: c_uint, ch: u32, fg: u16, bg: u16);

        //pub fn tb_select_input_mode(mode: c_int) -> c_int;
        //pub fn tb_set_clear_attributes(fg: u16, bg: u16);

        pub fn tb_peek_event(ev: *const ::RawEvent, timeout: c_uint) -> c_int;
        pub fn tb_poll_event(ev: *const ::RawEvent) -> c_int;
    }
}

fn nil_raw_event() -> RawEvent {
    RawEvent{etype: 0, emod: 0, key: 0, ch: 0, w: 0, h: 0}
}

fn unpack_event(ev_type: c_int, ev: &RawEvent) -> Event {
    match ev_type {
        0 => NoEvent,
        1 => {
            return KeyEvent(ev.emod, ev.key, ev.ch);
        },
        2 => {
            return ResizeEvent(ev.w, ev.h);
        },
        _ => { panic!("Unknown event"); }
    }
}

pub fn convert_color(c: Color) -> u16 {
    match c {
        Black   => 0x00,
        Red     => 0x01,
        Green   => 0x02,
        Yellow  => 0x03,
        Blue    => 0x04,
        Magenta => 0x05,
        Cyan    => 0x06,
        White   => 0x07,
    }
}

pub fn convert_style(sty: Style) -> u16 {
    match sty {
        Normal         => 0x00,
        Bold           => 0x10,
        Underline      => 0x20,
        BoldUnderline => 0x30,
    }
}

pub fn init() -> int {
    unsafe {
        termbox::tb_init() as int
    }
}

pub fn shutdown() {
    unsafe { termbox::tb_shutdown() }
}

pub fn width() -> uint {
    unsafe { termbox::tb_width() as uint }
}

pub fn height() -> uint {
    unsafe { termbox::tb_height() as uint }
}

pub fn clear() {
    unsafe { termbox::tb_clear() }
}

pub fn present() {
    unsafe { termbox::tb_present() }
}

pub fn set_cursor(x: int, y: int) {
    unsafe { termbox::tb_set_cursor(x as c_int, y as c_int) }
}

pub fn change_cell(x: uint, y: uint, ch: u32, fg: u16, bg: u16) {
    unsafe {
        termbox::tb_change_cell(x as c_uint, y as c_uint, ch, fg, bg)
    }
}

pub fn print(x: uint, y: uint, sty: Style, fg: Color, bg: Color, s: String) {
    let fg: u16 = convert_color(fg) | convert_style(sty);
    let bg: u16 = convert_color(bg);
    for (i, ch) in s.as_slice().chars().enumerate() {
        change_cell(x+i, y, ch as u32, fg, bg);
    }
}

pub fn print_char(x: uint, y: uint, sty: Style, fg: Color, bg: Color, ch: char) {
    let fg: u16 = convert_color(fg) | convert_style(sty);
    let bg: u16 = convert_color(bg);
    change_cell(x, y, ch as u32, fg, bg);
}

pub fn poll_event() -> Event {
    unsafe {
        let ev = nil_raw_event();
        let rc = termbox::tb_poll_event(&ev as *const RawEvent);
        return unpack_event(rc, &ev);
    }
}

pub fn peek_event(timeout: uint) -> Event {
    unsafe {
        let ev = nil_raw_event();
        let rc = termbox::tb_peek_event(&ev as *const RawEvent, timeout as c_uint);
        return unpack_event(rc, &ev);
    }
}
<|endoftext|>


<commit_before>extern crate html5ever;
extern crate tendril;
extern crate regex;
extern crate hyper;
extern crate string_cache;
extern crate url;

use html5ever::rcdom::{Document, Doctype, Text, Comment, Element};
use html5ever::rcdom::{RcDom, Handle};
use html5ever::{parse_document, Attribute};
use tendril::stream::TendrilSink;
use std::default::Default;

use regex::Regex;
use hyper::Client;
use hyper::header::Connection;
use hyper::header::ConnectionOption;

use Provider;
use Track;
use open_graph;
use soundcloud;
use youtube;
use error::Error;

use url::percent_encoding::{percent_decode};

static YOUTUBE_EMBED:       &'static str = r"www.youtube.com/embed/([a-zA-Z0-9_-].+)";
static YOUTUBE_LIST:        &'static str = r"www.youtube.com/embed/videoseries\?list=([a-zA-Z0-9_-]+)";
static YOUTUBE_WATCH:       &'static str = r"www.youtube.com/watch\?v=([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_TRACK:    &'static str = r"api.soundcloud.com/tracks/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_PLAYLIST: &'static str = r"api.soundcloud.com/playlists/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_USER:     &'static str = r"api.soundcloud.com/users/([a-zA-Z0-9_-]+)";

#[derive(Debug)]
pub struct ScraperProduct {
    pub tracks: Vec<Track>,
    pub og_obj: Option<open_graph::Object>,
}

pub fn extract(url: &str) -> Result<ScraperProduct, Error> {
    let client = Client::new();
    let mut res = client.get(url)
        .header(Connection(vec![ConnectionOption::Close]))
        .send()
        .unwrap();
    if res.status.is_success() {
        let dom = parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut res)
            .unwrap();
        let mut tracks   = Vec::new();
        let mut og_props = Vec::new();
        walk(0, dom.document, &mut tracks, &mut og_props);
        let og_obj = if og_props.len() > 0 {
            Some(open_graph::Object::new(&og_props))
        } else {
            None
        };
        Ok(ScraperProduct {
            tracks: tracks,
            og_obj: og_obj
        })
    } else {
        println!("Failed to get entry html {}: {}", res.status, url);
        Err(Error::NotFound)
    }
}

// This is not proper HTML serialization, of course.
fn walk(indent: usize,
        handle: Handle,
        tracks: &mut Vec<Track>,
        og_props: &mut Vec<(String, String)>) {
    let node = handle.borrow();
    match node.node {
        Document         => (),
        Doctype(_, _, _) => (),
        Text(_)          => (),
        Comment(_)       => (),
        Element(ref name, _, ref attrs) => {
            let tag_name = name.local.as_ref();
            let mut ps = extract_open_graph_metadata_from_tag(tag_name, attrs);
            og_props.append(&mut ps);
            let ts = extract_tracks_from_tag(tag_name, attrs);
            for track in ts.iter().cloned() {
                if !(tracks).iter().any(|t| track == *t) {
                    (*tracks).push(track)
                }
            }
        }
    }
    for child in node.children.iter() {
        walk(indent+4, child.clone(), tracks, og_props);
    }
}

fn attr(attr_name: &str, attrs: &Vec<Attribute>) -> Option<String> {
    for attr in attrs.iter() {
        if attr.name.local.as_ref() == attr_name {
            return Some(attr.value.to_string())
        }
    }
    None
}

pub fn extract_tracks_from_tag(tag_name: &str,
                               attrs: &Vec<Attribute>) -> Vec<Track> {
    if tag_name == "iframe" {
        match attr("src", attrs) {
            Some(ref src) => extract_tracks_from_url(src.to_string()),
            None => vec![]
        }
    } else if tag_name == "a" || tag_name == "link" {
        match attr("href", attrs) {
            Some(ref href) => extract_tracks_from_url(href.to_string()),
            None => vec![]
        }
    } else {
        vec![]
    }
}

pub fn extract_open_graph_metadata_from_tag(tag_name: &str,
                                            attrs: &Vec<Attribute>) -> Vec<(String, String)> {

    let mut og_props = vec!();
    if tag_name == "meta" {
        match extract_open_graph_prop("property", attrs) {
            Some((key, content)) => og_props.push((key, content)),
            None                 => (),
        }
        match extract_open_graph_prop("name", attrs) {
            Some((key, content)) => og_props.push((key, content)),
            None                 => (),
        }
    }
    og_props
}

fn extract_open_graph_prop<'a>(attr_name: &str, attrs: &Vec<Attribute>) -> Option<(String, String)> {
    attr(attr_name, attrs)
        .and_then(|property|
                  if property.starts_with("og:") {
                      let end = property.chars().count();
                      let key = unsafe {
                          property.slice_unchecked(3, end)
                      }.to_string();
                      attr("content", attrs).map(|content| (key, content))
                  } else {
                      None
                  })
}

fn extract_identifier(value: &str, regex_str: &str) -> Option<String> {
    match Regex::new(regex_str) {
        Ok(re) => match re.captures(value) {
            Some(cap) => match cap.at(1) {
                Some(str) => {
                    let strs: Vec<&str> = str.split('?').collect();
                    return Some(strs[0].to_string())
                },
                None => None
            },
            None => None
        },
        Err(_) => None
    }
}

fn extract_tracks_from_url(url: String) -> Vec<Track> {
    let decoded = percent_decode(url.as_bytes()).decode_utf8_lossy().into_owned();
    match extract_identifier(&decoded, YOUTUBE_WATCH) {
        Some(identifier) => {
            return vec![Track::new(Provider::YouTube, identifier)]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_LIST) {
        Some(identifier) => {
            return match youtube::fetch_playlist(&identifier) {
                Ok(res) => res.items.iter()
                                    .map(|ref i| Track::from_yt_playlist_item(i))
                                    .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_EMBED) {
        Some(identifier) => {
            return vec![Track::new(Provider::YouTube, identifier)]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_TRACK) {
        Some(identifier) => {
            return vec![Track::new(Provider::SoundCloud, identifier)]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_PLAYLIST) {
        Some(identifier) => {
            return match soundcloud::fetch_playlist(&identifier) {
                Ok(playlist) => playlist.tracks
                                        .iter()
                                        .map(|ref t| Track::from_sc_track(t))
                                        .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_USER) {
        Some(identifier) => {
            print!("user {}", identifier);
            return match soundcloud::fetch_user_tracks(&identifier) {
                Ok(tracks) => tracks.iter()
                                    .map(|ref t| Track::from_sc_track(t))
                                    .collect::<Vec<_>>(),
                Err(_)     => vec![]
            }
        },
        None => ()
    }
    return vec![]
}

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

    #[test]
    fn test_extract_identifier() {
        let soundcloud_src = "https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/tracks/195425494/stream&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&visual=true";
        match extract_identifier(soundcloud_src, super::SOUNDCLOUD_TRACK) {
            Some(identifier) => assert_eq!(identifier, "195425494".to_string()),
            None             => assert!(false)
        }
        let youtube_embed = "https://www.youtube.com/embed/X8tOngmlES0?rel=0";
        match extract_identifier(youtube_embed, super::YOUTUBE_EMBED) {
            Some(identifier) => assert_eq!(identifier, "X8tOngmlES0".to_string()),
            None             => assert!(false)
        }

        let youtube_watch = "https://www.youtube.com/watch?v=oDuif301F-8";
        match extract_identifier(youtube_watch, super::YOUTUBE_WATCH) {
            Some(identifier) => assert_eq!(identifier, "oDuif301F-8".to_string()),
            None             => assert!(false)
        }

        let youtube_list = "https://www.youtube.com/embed/videoseries?list=PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX";
        match extract_identifier(youtube_list, super::YOUTUBE_LIST) {
            Some(identifier) => assert_eq!(identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string()),
            None             => assert!(false)
        }
    }
}
<commit_msg>Adopt new regex crate<commit_after>extern crate html5ever;
extern crate tendril;
extern crate regex;
extern crate hyper;
extern crate string_cache;
extern crate url;

use html5ever::rcdom::{Document, Doctype, Text, Comment, Element};
use html5ever::rcdom::{RcDom, Handle};
use html5ever::{parse_document, Attribute};
use tendril::stream::TendrilSink;
use std::default::Default;

use regex::Regex;
use hyper::Client;
use hyper::header::Connection;
use hyper::header::ConnectionOption;

use Provider;
use Track;
use open_graph;
use soundcloud;
use youtube;
use error::Error;

use url::percent_encoding::{percent_decode};

static YOUTUBE_EMBED:       &'static str = r"www.youtube.com/embed/([a-zA-Z0-9_-].+)";
static YOUTUBE_LIST:        &'static str = r"www.youtube.com/embed/videoseries\?list=([a-zA-Z0-9_-]+)";
static YOUTUBE_WATCH:       &'static str = r"www.youtube.com/watch\?v=([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_TRACK:    &'static str = r"api.soundcloud.com/tracks/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_PLAYLIST: &'static str = r"api.soundcloud.com/playlists/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_USER:     &'static str = r"api.soundcloud.com/users/([a-zA-Z0-9_-]+)";

#[derive(Debug)]
pub struct ScraperProduct {
    pub tracks: Vec<Track>,
    pub og_obj: Option<open_graph::Object>,
}

pub fn extract(url: &str) -> Result<ScraperProduct, Error> {
    let client = Client::new();
    let mut res = client.get(url)
        .header(Connection(vec![ConnectionOption::Close]))
        .send()
        .unwrap();
    if res.status.is_success() {
        let dom = parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut res)
            .unwrap();
        let mut tracks   = Vec::new();
        let mut og_props = Vec::new();
        walk(0, dom.document, &mut tracks, &mut og_props);
        let og_obj = if og_props.len() > 0 {
            Some(open_graph::Object::new(&og_props))
        } else {
            None
        };
        Ok(ScraperProduct {
            tracks: tracks,
            og_obj: og_obj
        })
    } else {
        println!("Failed to get entry html {}: {}", res.status, url);
        Err(Error::NotFound)
    }
}

// This is not proper HTML serialization, of course.
fn walk(indent: usize,
        handle: Handle,
        tracks: &mut Vec<Track>,
        og_props: &mut Vec<(String, String)>) {
    let node = handle.borrow();
    match node.node {
        Document         => (),
        Doctype(_, _, _) => (),
        Text(_)          => (),
        Comment(_)       => (),
        Element(ref name, _, ref attrs) => {
            let tag_name = name.local.as_ref();
            let mut ps = extract_open_graph_metadata_from_tag(tag_name, attrs);
            og_props.append(&mut ps);
            let ts = extract_tracks_from_tag(tag_name, attrs);
            for track in ts.iter().cloned() {
                if !(tracks).iter().any(|t| track == *t) {
                    (*tracks).push(track)
                }
            }
        }
    }
    for child in node.children.iter() {
        walk(indent+4, child.clone(), tracks, og_props);
    }
}

fn attr(attr_name: &str, attrs: &Vec<Attribute>) -> Option<String> {
    for attr in attrs.iter() {
        if attr.name.local.as_ref() == attr_name {
            return Some(attr.value.to_string())
        }
    }
    None
}

pub fn extract_tracks_from_tag(tag_name: &str,
                               attrs: &Vec<Attribute>) -> Vec<Track> {
    if tag_name == "iframe" {
        match attr("src", attrs) {
            Some(ref src) => extract_tracks_from_url(src.to_string()),
            None => vec![]
        }
    } else if tag_name == "a" || tag_name == "link" {
        match attr("href", attrs) {
            Some(ref href) => extract_tracks_from_url(href.to_string()),
            None => vec![]
        }
    } else {
        vec![]
    }
}

pub fn extract_open_graph_metadata_from_tag(tag_name: &str,
                                            attrs: &Vec<Attribute>) -> Vec<(String, String)> {

    let mut og_props = vec!();
    if tag_name == "meta" {
        match extract_open_graph_prop("property", attrs) {
            Some((key, content)) => og_props.push((key, content)),
            None                 => (),
        }
        match extract_open_graph_prop("name", attrs) {
            Some((key, content)) => og_props.push((key, content)),
            None                 => (),
        }
    }
    og_props
}

fn extract_open_graph_prop<'a>(attr_name: &str, attrs: &Vec<Attribute>) -> Option<(String, String)> {
    attr(attr_name, attrs)
        .and_then(|property|
                  if property.starts_with("og:") {
                      let end = property.chars().count();
                      let key = unsafe {
                          property.slice_unchecked(3, end)
                      }.to_string();
                      attr("content", attrs).map(|content| (key, content))
                  } else {
                      None
                  })
}

fn extract_identifier(value: &str, regex_str: &str) -> Option<String> {
    match Regex::new(regex_str) {
        Ok(re) => match re.captures(value) {
            Some(cap) => {
                let strs: Vec<&str> = cap[1].split('?').collect();
                return Some(strs[0].to_string())
            },
            None => None
        },
        Err(_) => None
    }
}

fn extract_tracks_from_url(url: String) -> Vec<Track> {
    let decoded = percent_decode(url.as_bytes()).decode_utf8_lossy().into_owned();
    match extract_identifier(&decoded, YOUTUBE_WATCH) {
        Some(identifier) => {
            return vec![Track::new(Provider::YouTube, identifier)]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_LIST) {
        Some(identifier) => {
            return match youtube::fetch_playlist(&identifier) {
                Ok(res) => res.items.iter()
                                    .map(|ref i| Track::from_yt_playlist_item(i))
                                    .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_EMBED) {
        Some(identifier) => {
            return vec![Track::new(Provider::YouTube, identifier)]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_TRACK) {
        Some(identifier) => {
            return vec![Track::new(Provider::SoundCloud, identifier)]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_PLAYLIST) {
        Some(identifier) => {
            return match soundcloud::fetch_playlist(&identifier) {
                Ok(playlist) => playlist.tracks
                                        .iter()
                                        .map(|ref t| Track::from_sc_track(t))
                                        .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_USER) {
        Some(identifier) => {
            print!("user {}", identifier);
            return match soundcloud::fetch_user_tracks(&identifier) {
                Ok(tracks) => tracks.iter()
                                    .map(|ref t| Track::from_sc_track(t))
                                    .collect::<Vec<_>>(),
                Err(_)     => vec![]
            }
        },
        None => ()
    }
    return vec![]
}

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

    #[test]
    fn test_extract_identifier() {
        let soundcloud_src = "https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/tracks/195425494/stream&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&visual=true";
        match extract_identifier(soundcloud_src, super::SOUNDCLOUD_TRACK) {
            Some(identifier) => assert_eq!(identifier, "195425494".to_string()),
            None             => assert!(false)
        }
        let youtube_embed = "https://www.youtube.com/embed/X8tOngmlES0?rel=0";
        match extract_identifier(youtube_embed, super::YOUTUBE_EMBED) {
            Some(identifier) => assert_eq!(identifier, "X8tOngmlES0".to_string()),
            None             => assert!(false)
        }

        let youtube_watch = "https://www.youtube.com/watch?v=oDuif301F-8";
        match extract_identifier(youtube_watch, super::YOUTUBE_WATCH) {
            Some(identifier) => assert_eq!(identifier, "oDuif301F-8".to_string()),
            None             => assert!(false)
        }

        let youtube_list = "https://www.youtube.com/embed/videoseries?list=PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX";
        match extract_identifier(youtube_list, super::YOUTUBE_LIST) {
            Some(identifier) => assert_eq!(identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string()),
            None             => assert!(false)
        }
    }
}
<|endoftext|>