<commit_before>
use objc;

use cocoa::base::{id, nil, class, NO};
use cocoa::foundation::{NSUInteger, NSRect, NSPoint, NSSize, NSAutoreleasePool, NSString};

use cocoa::appkit::{NSWindow, NSBackingStoreBuffered};
use cocoa::appkit::{NSViewHeightSizable, NSViewWidthSizable};
use cocoa::appkit::{NSTitledWindowMask, NSMiniaturizableWindowMask, NSResizableWindowMask};

use std::sync::atomic::AtomicPtr;
use std::sync::Arc;

use super::view::View;
use atomic_box::AtomicBox;

#[repr(u64)]
enum UserInterfaceLayoutOrientation {
    Vertical = 1,
}

pub struct Window<M> {
    _id: AtomicPtr<objc::runtime::Object>,
    stack: AtomicPtr<objc::runtime::Object>,
    views: Arc<Vec<AtomicBox<Box<View<M>>>>>,
}

impl<M> Window<M> {
    pub fn new(title: &str) -> Self {
        unsafe {
            let style = NSResizableWindowMask as NSUInteger | NSTitledWindowMask as NSUInteger |
                        NSMiniaturizableWindowMask as NSUInteger;
            let rect = NSRect::new(NSPoint::new(0., 0.), NSSize::new(350., 350.));
            let window = NSWindow::alloc(nil)
                .initWithContentRect_styleMask_backing_defer_(rect,
                                                              style,
                                                              NSBackingStoreBuffered,
                                                              NO)
                .autorelease();
            window.cascadeTopLeftFromPoint_(NSPoint::new(20., 20.));
            window.center();

            let title = NSString::alloc(nil).init_str(title);
            window.setTitle_(title);

            window.makeKeyAndOrderFront_(nil);

            let stack = {
                let stack: id = msg_send![class("NSStackView"), alloc];
                let stack: id = msg_send![stack, init];

                // window.frame padded by 10.0 on each side
                let rect = NSRect::new(NSPoint::new(10., 10.), NSSize::new(330., 330.));
                msg_send![stack, setFrame: rect];

                msg_send![stack, setAutoresizingMask: NSViewWidthSizable | NSViewHeightSizable];

                msg_send![stack, setOrientation: UserInterfaceLayoutOrientation::Vertical];
                msg_send![stack, setSpacing: 25.0];

                use cocoa::appkit::NSView;
                window.contentView().addSubview_(stack);

                stack
            };

            Window {
                _id: AtomicPtr::new(window),
                stack: AtomicPtr::new(stack),
                views: Arc::new(Vec::new()),
            }
        }
    }

    pub fn add<V: View<M> + 'static>(&mut self, view: V) {
        unsafe { msg_send![*self.stack.get_mut(), addArrangedSubview: view.id()] };

        if let Some(views) = Arc::get_mut(&mut self.views) {
            views.push(AtomicBox::new(Box::new(Box::new(view))));
        }
    }

    pub fn update(&mut self, model: &M) {
        if let Some(views) = Arc::get_mut(&mut self.views) {
            for view in views.iter_mut() {
                view.get_mut().update(model);
            }
        }
    }
}
<commit_msg>Use addView:inGravity: instead of addArrangedSubview: (#8)<commit_after>
use objc;

use cocoa::base::{id, nil, class, NO};
use cocoa::foundation::{NSUInteger, NSRect, NSPoint, NSSize, NSAutoreleasePool, NSString};

use cocoa::appkit::{NSWindow, NSBackingStoreBuffered};
use cocoa::appkit::{NSViewHeightSizable, NSViewWidthSizable};
use cocoa::appkit::{NSTitledWindowMask, NSMiniaturizableWindowMask, NSResizableWindowMask};

use std::sync::atomic::AtomicPtr;
use std::sync::Arc;

use super::view::View;
use atomic_box::AtomicBox;

#[repr(u64)]
enum UserInterfaceLayoutOrientation {
    Vertical = 1,
}

#[repr(u64)]
enum NSStackViewGravity {
    Top = 1,
}

pub struct Window<M> {
    _id: AtomicPtr<objc::runtime::Object>,
    stack: AtomicPtr<objc::runtime::Object>,
    views: Arc<Vec<AtomicBox<Box<View<M>>>>>,
}

impl<M> Window<M> {
    pub fn new(title: &str) -> Self {
        unsafe {
            let style = NSResizableWindowMask as NSUInteger | NSTitledWindowMask as NSUInteger |
                        NSMiniaturizableWindowMask as NSUInteger;
            let rect = NSRect::new(NSPoint::new(0., 0.), NSSize::new(350., 350.));
            let window = NSWindow::alloc(nil)
                .initWithContentRect_styleMask_backing_defer_(rect,
                                                              style,
                                                              NSBackingStoreBuffered,
                                                              NO)
                .autorelease();
            window.cascadeTopLeftFromPoint_(NSPoint::new(20., 20.));
            window.center();

            let title = NSString::alloc(nil).init_str(title);
            window.setTitle_(title);

            window.makeKeyAndOrderFront_(nil);

            let stack = {
                let stack: id = msg_send![class("NSStackView"), alloc];
                let stack: id = msg_send![stack, init];

                // window.frame padded by 10.0 on each side
                let rect = NSRect::new(NSPoint::new(10., 10.), NSSize::new(330., 330.));
                msg_send![stack, setFrame: rect];

                msg_send![stack, setAutoresizingMask: NSViewWidthSizable | NSViewHeightSizable];

                msg_send![stack, setOrientation: UserInterfaceLayoutOrientation::Vertical];
                msg_send![stack, setSpacing: 25.0];

                use cocoa::appkit::NSView;
                window.contentView().addSubview_(stack);

                stack
            };

            Window {
                _id: AtomicPtr::new(window),
                stack: AtomicPtr::new(stack),
                views: Arc::new(Vec::new()),
            }
        }
    }

    pub fn add<V: View<M> + 'static>(&mut self, view: V) {
        unsafe {
            msg_send![*self.stack.get_mut(), addView:view.id()
                                           inGravity:NSStackViewGravity::Top]
        };

        if let Some(views) = Arc::get_mut(&mut self.views) {
            views.push(AtomicBox::new(Box::new(Box::new(view))));
        }
    }

    pub fn update(&mut self, model: &M) {
        if let Some(views) = Arc::get_mut(&mut self.views) {
            for view in views.iter_mut() {
                view.get_mut().update(model);
            }
        }
    }
}
<|endoftext|>


<commit_before>use crate::consts::{miri_to_const, Constant};
use crate::utils::{in_macro, is_type_diagnostic_item, snippet, span_lint_and_sugg, walk_ptrs_ty};
use if_chain::if_chain;
use rustc_ast::ast::LitKind;
use rustc_errors::Applicability;
use rustc_hir::def::{DefKind, Res};
use rustc_hir::{Expr, ExprKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty::{self, Ty};
use rustc_session::{declare_lint_pass, declare_tool_lint};

declare_clippy_lint! {
    /// **What it does:** Checks for usage of `.repeat(1)` and suggest the following method for each types.
    /// - `.to_string()` for `str`
    /// - `.clone()` for `String`
    /// - `.to_vec()` for `slice`
    ///
    /// **Why is this bad?** For example, `String.repeat(1)` is equivalent to `.clone()`. If cloning the string is the intention behind this, `clone()` should be used.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    ///
    /// ```rust
    /// fn main() {
    ///     let x = String::from("hello world").repeat(1);
    /// }
    /// ```
    /// Use instead:
    /// ```rust
    /// fn main() {
    ///     let x = String::from("hello world").clone();
    /// }
    /// ```
    pub REPEAT_ONCE,
    complexity,
    "using `.repeat(1)` instead of `String.clone()`, `str.to_string()` or `slice.to_vec()` "
}

declare_lint_pass!(RepeatOnce => [REPEAT_ONCE]);

impl<'tcx> LateLintPass<'tcx> for RepeatOnce {
    fn check_expr(&mut self, cx: &LateContext<'_>, expr: &'tcx Expr<'_>) {
        if_chain! {
            if let ExprKind::MethodCall(ref path, _, ref args, _) = expr.kind;
            if path.ident.name == sym!(repeat);
            if is_once(cx, &args[1]) && !in_macro(args[0].span);
            then {
                let ty = walk_ptrs_ty(cx.tables().expr_ty(&args[0]));
                if is_str(ty){
                    span_lint_and_sugg(
                        cx,
                        REPEAT_ONCE,
                        expr.span,
                        "calling `repeat(1)` on str",
                        "consider using `.to_string()` instead",
                        format!("{}.to_string()", snippet(cx, args[0].span, r#""...""#)),
                        Applicability::MachineApplicable,
                    );
                } else if is_slice(ty) {
                    span_lint_and_sugg(
                        cx,
                        REPEAT_ONCE,
                        expr.span,
                        "calling `repeat(1)` on slice",
                        "consider using `.to_vec()` instead",
                        format!("{}.to_vec()", snippet(cx, args[0].span, r#""...""#)),
                        Applicability::MachineApplicable,
                    );
                } else if is_type_diagnostic_item(cx, ty, sym!(string_type)) {
                    span_lint_and_sugg(
                        cx,
                        REPEAT_ONCE,
                        expr.span,
                        "calling `repeat(1)` on a string literal",
                        "consider using `.clone()` instead",
                        format!("{}.clone()", snippet(cx, args[0].span, r#""...""#)),
                        Applicability::MachineApplicable,
                    );
                }
            }
        }
    }
}

fn is_once<'tcx>(cx: &LateContext<'_>, expr: &'tcx Expr<'_>) -> bool {
    match expr.kind {
        ExprKind::Lit(ref lit) => {
            if let LitKind::Int(ref lit_content, _) = lit.node {
                *lit_content == 1
            } else {
                false
            }
        },
        ExprKind::Path(rustc_hir::QPath::Resolved(None, path)) => {
            if let Res::Def(DefKind::Const, def_id) = path.res {
                let ty = cx.tcx.type_of(def_id);
                let con = cx
                    .tcx
                    .const_eval_poly(def_id)
                    .ok()
                    .map(|val| rustc_middle::ty::Const::from_value(cx.tcx, val, ty))
                    .unwrap();
                let con = miri_to_const(con);
                con == Some(Constant::Int(1))
            } else {
                false
            }
        },
        _ => false,
    }
}

fn is_str(ty: Ty<'_>) -> bool {
    match ty.kind {
        ty::Str => true,
        _ => false,
    }
}

fn is_slice(ty: Ty<'_>) -> bool {
    match ty.kind {
        ty::Slice(..) | ty::Array(..) => true,
        _ => false,
    }
}
<commit_msg>Refactoring to use `constant_context<commit_after>use crate::consts::{constant_context, Constant};
use crate::utils::{in_macro, is_type_diagnostic_item, snippet, span_lint_and_sugg, walk_ptrs_ty};
use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::{Expr, ExprKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_session::{declare_lint_pass, declare_tool_lint};

declare_clippy_lint! {
    /// **What it does:** Checks for usage of `.repeat(1)` and suggest the following method for each types.
    /// - `.to_string()` for `str`
    /// - `.clone()` for `String`
    /// - `.to_vec()` for `slice`
    ///
    /// **Why is this bad?** For example, `String.repeat(1)` is equivalent to `.clone()`. If cloning the string is the intention behind this, `clone()` should be used.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    ///
    /// ```rust
    /// fn main() {
    ///     let x = String::from("hello world").repeat(1);
    /// }
    /// ```
    /// Use instead:
    /// ```rust
    /// fn main() {
    ///     let x = String::from("hello world").clone();
    /// }
    /// ```
    pub REPEAT_ONCE,
    complexity,
    "using `.repeat(1)` instead of `String.clone()`, `str.to_string()` or `slice.to_vec()` "
}

declare_lint_pass!(RepeatOnce => [REPEAT_ONCE]);

impl<'tcx> LateLintPass<'tcx> for RepeatOnce {
    fn check_expr(&mut self, cx: &LateContext<'_>, expr: &'tcx Expr<'_>) {
        if_chain! {
            if let ExprKind::MethodCall(ref path, _, ref args, _) = expr.kind;
            if path.ident.name == sym!(repeat);
            if let Some(Constant::Int(1)) = constant_context(cx, cx.tables()).expr(&args[1]);
            if !in_macro(args[0].span);
            then {
                let ty = walk_ptrs_ty(cx.tables().expr_ty(&args[0]));
                if ty.is_str() {
                    span_lint_and_sugg(
                        cx,
                        REPEAT_ONCE,
                        expr.span,
                        "calling `repeat(1)` on str",
                        "consider using `.to_string()` instead",
                        format!("{}.to_string()", snippet(cx, args[0].span, r#""...""#)),
                        Applicability::MachineApplicable,
                    );
                } else if let Some(_) = ty.builtin_index() {
                    span_lint_and_sugg(
                        cx,
                        REPEAT_ONCE,
                        expr.span,
                        "calling `repeat(1)` on slice",
                        "consider using `.to_vec()` instead",
                        format!("{}.to_vec()", snippet(cx, args[0].span, r#""...""#)),
                        Applicability::MachineApplicable,
                    );
                } else if is_type_diagnostic_item(cx, ty, sym!(string_type)) {
                    span_lint_and_sugg(
                        cx,
                        REPEAT_ONCE,
                        expr.span,
                        "calling `repeat(1)` on a string literal",
                        "consider using `.clone()` instead",
                        format!("{}.clone()", snippet(cx, args[0].span, r#""...""#)),
                        Applicability::MachineApplicable,
                    );
                }
            }
        }
    }
}
<|endoftext|>


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

use std::{
    collections::HashSet,
    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 shadowsocks::{context::Context, relay::socks5::Address};

use self::sub_domains_tree::SubDomainsTree;

mod sub_domains_tree;

/// 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_regex: RegexSet,
    rule_set: HashSet<String>,
    rule_tree: SubDomainsTree,
}

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

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

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

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

        write!(f, "], rule_set: {:?}, rule_tree: {:?} }}", self.rule_set, self.rule_tree)
    }
}

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

        Rules {
            ipv4,
            ipv6,
            rule_regex,
            rule_set,
            rule_tree,
        }
    }

    /// Check if the specified address matches these rules
    #[allow(dead_code)]
    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_set.contains(host) || self.rule_tree.contains(host) || self.rule_regex.is_match(host)
    }

    /// Check if there are no rules for IP addresses
    fn is_ip_empty(&self) -> bool {
        self.ipv4.is_empty() && self.ipv6.is_empty()
    }

    /// Check if there are no rules for domain names
    fn is_host_empty(&self) -> bool {
        self.rule_set.is_empty() && self.rule_tree.is_empty() && self.rule_regex.is_empty()
    }
}

/// 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_regex = Vec::new();
        let mut outbound_block_rules_set = HashSet::new();
        let mut outbound_block_rules_tree = SubDomainsTree::new();
        let mut bypass_ipv4 = IpRange::new();
        let mut bypass_ipv6 = IpRange::new();
        let mut bypass_rules_regex = Vec::new();
        let mut bypass_rules_set = HashSet::new();
        let mut bypass_rules_tree = SubDomainsTree::new();
        let mut proxy_ipv4 = IpRange::new();
        let mut proxy_ipv6 = IpRange::new();
        let mut proxy_rules_regex = Vec::new();
        let mut proxy_rules_set = HashSet::new();
        let mut proxy_rules_tree = SubDomainsTree::new();

        let mut curr_ipv4 = &mut bypass_ipv4;
        let mut curr_ipv6 = &mut bypass_ipv6;
        let mut curr_rules_regex = &mut bypass_rules_regex;
        let mut curr_rules_set = &mut bypass_rules_set;
        let mut curr_rules_tree = &mut bypass_rules_tree;

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

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

            if line.starts_with("||") {
                curr_rules_tree.insert(&line[2..]);
                continue;
            }

            if line.starts_with('|') {
                curr_rules_set.insert(line[1..].to_string());
                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_regex = &mut outbound_block_rules_regex;
                    curr_rules_set = &mut outbound_block_rules_set;
                    curr_rules_tree = &mut outbound_block_rules_tree;
                }
                "[black_list]" | "[bypass_list]" => {
                    curr_ipv4 = &mut bypass_ipv4;
                    curr_ipv6 = &mut bypass_ipv6;
                    curr_rules_regex = &mut bypass_rules_regex;
                    curr_rules_set = &mut bypass_rules_set;
                    curr_rules_tree = &mut bypass_rules_tree;
                }
                "[white_list]" | "[proxy_list]" => {
                    curr_ipv4 = &mut proxy_ipv4;
                    curr_ipv6 = &mut proxy_ipv6;
                    curr_rules_regex = &mut proxy_rules_regex;
                    curr_rules_set = &mut proxy_rules_set;
                    curr_rules_tree = &mut proxy_rules_tree;
                }
                _ => {
                    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_regex.push(line);
                                }
                            }
                        }
                    }
                }
            }
        }

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

        let outbound_block_regex = match RegexSetBuilder::new(outbound_block_rules_regex)
            .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_regex)
            .case_insensitive(true)
            .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_regex)
            .case_insensitive(true)
            .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,
                outbound_block_rules_set,
                outbound_block_rules_tree,
            ),
            black_list: Rules::new(
                bypass_ipv4,
                bypass_ipv6,
                bypass_regex,
                bypass_rules_set,
                bypass_rules_tree,
            ),
            white_list: Rules::new(proxy_ipv4, proxy_ipv6, proxy_regex, proxy_rules_set, proxy_rules_tree),
            mode,
        })
    }

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    ///
    /// Return
    /// - `Some(true)` if `host` is in `white_list` (should be proxied)
    /// - `Some(false)` if `host` is in `black_list` (should be bypassed)
    /// - `None` if `host` doesn't match any rules
    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 IP rules
    pub fn is_ip_empty(&self) -> bool {
        match self.mode {
            Mode::BlackList => self.black_list.is_ip_empty(),
            Mode::WhiteList => self.white_list.is_ip_empty(),
        }
    }

    /// If there are no domain name rules
    pub fn is_host_empty(&self) -> bool {
        self.black_list.is_host_empty() && self.white_list.is_host_empty()
    }

    /// Check if `IpAddr` should be proxied
    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),
        }
    }

    /// Default mode
    ///
    /// Default behavor for hosts that are not configured
    /// - `true` - Proxied
    /// - `false` - Bypassed
    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_ip_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 async fn check_outbound_blocked(&self, context: &Context, outbound: &Address) -> bool {
        match outbound {
            Address::SocketAddress(saddr) => self.outbound_block.check_ip_matched(&saddr.ip()),
            Address::DomainNameAddress(host, port) => {
                if self.outbound_block.check_host_matched(host) {
                    return true;
                }

                if let Ok(vaddr) = context.dns_resolve(host, *port).await {
                    for addr in vaddr {
                        if self.outbound_block.check_ip_matched(&addr.ip()) {
                            return true;
                        }
                    }
                }

                false
            }
        }
    }
}
<commit_msg>Disable unicode matching in ACL rules<commit_after>//! Access Control List (ACL) for shadowsocks
//!
//! This is for advance controlling server behaviors in both local and proxy servers.

use std::{
    collections::HashSet,
    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 shadowsocks::{context::Context, relay::socks5::Address};

use self::sub_domains_tree::SubDomainsTree;

mod sub_domains_tree;

/// 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_regex: RegexSet,
    rule_set: HashSet<String>,
    rule_tree: SubDomainsTree,
}

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

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

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

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

        write!(f, "], rule_set: {:?}, rule_tree: {:?} }}", self.rule_set, self.rule_tree)
    }
}

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

        Rules {
            ipv4,
            ipv6,
            rule_regex,
            rule_set,
            rule_tree,
        }
    }

    /// Check if the specified address matches these rules
    #[allow(dead_code)]
    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_set.contains(host) || self.rule_tree.contains(host) || self.rule_regex.is_match(host)
    }

    /// Check if there are no rules for IP addresses
    fn is_ip_empty(&self) -> bool {
        self.ipv4.is_empty() && self.ipv6.is_empty()
    }

    /// Check if there are no rules for domain names
    fn is_host_empty(&self) -> bool {
        self.rule_set.is_empty() && self.rule_tree.is_empty() && self.rule_regex.is_empty()
    }
}

/// 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_regex = Vec::new();
        let mut outbound_block_rules_set = HashSet::new();
        let mut outbound_block_rules_tree = SubDomainsTree::new();
        let mut bypass_ipv4 = IpRange::new();
        let mut bypass_ipv6 = IpRange::new();
        let mut bypass_rules_regex = Vec::new();
        let mut bypass_rules_set = HashSet::new();
        let mut bypass_rules_tree = SubDomainsTree::new();
        let mut proxy_ipv4 = IpRange::new();
        let mut proxy_ipv6 = IpRange::new();
        let mut proxy_rules_regex = Vec::new();
        let mut proxy_rules_set = HashSet::new();
        let mut proxy_rules_tree = SubDomainsTree::new();

        let mut curr_ipv4 = &mut bypass_ipv4;
        let mut curr_ipv6 = &mut bypass_ipv6;
        let mut curr_rules_regex = &mut bypass_rules_regex;
        let mut curr_rules_set = &mut bypass_rules_set;
        let mut curr_rules_tree = &mut bypass_rules_tree;

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

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

            if line.starts_with("||") {
                curr_rules_tree.insert(&line[2..]);
                continue;
            }

            if line.starts_with('|') {
                curr_rules_set.insert(line[1..].to_string());
                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_regex = &mut outbound_block_rules_regex;
                    curr_rules_set = &mut outbound_block_rules_set;
                    curr_rules_tree = &mut outbound_block_rules_tree;
                }
                "[black_list]" | "[bypass_list]" => {
                    curr_ipv4 = &mut bypass_ipv4;
                    curr_ipv6 = &mut bypass_ipv6;
                    curr_rules_regex = &mut bypass_rules_regex;
                    curr_rules_set = &mut bypass_rules_set;
                    curr_rules_tree = &mut bypass_rules_tree;
                }
                "[white_list]" | "[proxy_list]" => {
                    curr_ipv4 = &mut proxy_ipv4;
                    curr_ipv6 = &mut proxy_ipv6;
                    curr_rules_regex = &mut proxy_rules_regex;
                    curr_rules_set = &mut proxy_rules_set;
                    curr_rules_tree = &mut proxy_rules_tree;
                }
                _ => {
                    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_regex.push(line);
                                }
                            }
                        }
                    }
                }
            }
        }

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

        let outbound_block_regex = match RegexSetBuilder::new(outbound_block_rules_regex)
            .size_limit(REGEX_SIZE_LIMIT)
            .unicode(false)
            .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_regex)
            .case_insensitive(true)
            .size_limit(REGEX_SIZE_LIMIT)
            .unicode(false)
            .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_regex)
            .case_insensitive(true)
            .size_limit(REGEX_SIZE_LIMIT)
            .unicode(false)
            .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,
                outbound_block_rules_set,
                outbound_block_rules_tree,
            ),
            black_list: Rules::new(
                bypass_ipv4,
                bypass_ipv6,
                bypass_regex,
                bypass_rules_set,
                bypass_rules_tree,
            ),
            white_list: Rules::new(proxy_ipv4, proxy_ipv6, proxy_regex, proxy_rules_set, proxy_rules_tree),
            mode,
        })
    }

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    ///
    /// Return
    /// - `Some(true)` if `host` is in `white_list` (should be proxied)
    /// - `Some(false)` if `host` is in `black_list` (should be bypassed)
    /// - `None` if `host` doesn't match any rules
    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 IP rules
    pub fn is_ip_empty(&self) -> bool {
        match self.mode {
            Mode::BlackList => self.black_list.is_ip_empty(),
            Mode::WhiteList => self.white_list.is_ip_empty(),
        }
    }

    /// If there are no domain name rules
    pub fn is_host_empty(&self) -> bool {
        self.black_list.is_host_empty() && self.white_list.is_host_empty()
    }

    /// Check if `IpAddr` should be proxied
    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),
        }
    }

    /// Default mode
    ///
    /// Default behavor for hosts that are not configured
    /// - `true` - Proxied
    /// - `false` - Bypassed
    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_ip_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 async fn check_outbound_blocked(&self, context: &Context, outbound: &Address) -> bool {
        match outbound {
            Address::SocketAddress(saddr) => self.outbound_block.check_ip_matched(&saddr.ip()),
            Address::DomainNameAddress(host, port) => {
                if self.outbound_block.check_host_matched(host) {
                    return true;
                }

                if let Ok(vaddr) = context.dns_resolve(host, *port).await {
                    for addr in vaddr {
                        if self.outbound_block.check_ip_matched(&addr.ip()) {
                            return true;
                        }
                    }
                }

                false
            }
        }
    }
}
<|endoftext|>


<commit_before>use anyhow::{bail, Result};

use crate::task_handler::ProcessAction;

/// Send a signal to a windows process.
pub fn send_signal(pid: u32, action: &ProcessAction, children: bool) -> Result<bool> {
    bail!("not supported on windows.")
}

/// Get all children of a specific process.
/// The Vec<i32> is just a placeholder.
pub fn get_children(pid: i32) -> Option<Vec<i32>> {
    None
}

/// Send a signal to multiple processes.
/// The Vec<i32> is just a placeholder.
pub fn send_signal_to_processes(processes: Vec<i32>, action: &ProcessAction) {}
<commit_msg>Fix windows clippy issues<commit_after>use anyhow::{bail, Result};

use crate::task_handler::ProcessAction;

/// Send a signal to a windows process.
pub fn send_signal(_pid: u32, _action: &ProcessAction, _children: bool) -> Result<bool> {
    bail!("not supported on windows.")
}

/// Get all children of a specific process.
/// The Vec<i32> is just a placeholder.
pub fn get_children(_pid: i32) -> Option<Vec<i32>> {
    None
}

/// Send a signal to multiple processes.
/// The Vec<i32> is just a placeholder.
pub fn send_signal_to_processes(_processes: Vec<i32>, _action: &ProcessAction) {}
<|endoftext|>


<commit_before>//! Base trait and common functionalities for all consumers.
extern crate rdkafka_sys as rdkafka;

pub mod base_consumer;
pub mod stream_consumer;

use self::rdkafka::types::*;

use std::ptr;
use std::os::raw::c_void;
use util::cstr_to_owned;

use client::{Context, NativeClient};
use message::Message;
use metadata::Metadata;
use error::KafkaResult;

pub use consumer::base_consumer::BaseConsumer;
pub use topic_partition_list::TopicPartitionList;

/// Rebalance information.
#[derive(Clone, Debug)]
pub enum Rebalance {
    Assign(TopicPartitionList),
    Revoke,
    Error(String)
}

/// Consumer specific Context. This user-defined object can be used to provide custom callbacks to
/// consumer events. Refer to the list of methods to check which callbacks can be specified.
pub trait ConsumerContext: Context {
    /// Implements the default rebalancing strategy and calls the `pre_rebalance` and
    /// `post_rebalance` methods. If this method is overridden, it will be responsibility
    /// of the user to call them if needed.
    fn rebalance(&self, native_client: &NativeClient, err: RDKafkaRespErr,
                 partitions_ptr: *mut RDKafkaTopicPartitionList) {

        let rebalance = match err {
            RDKafkaRespErr::RD_KAFKA_RESP_ERR__ASSIGN_PARTITIONS => {
                // TODO: this might be expensive
                let topic_partition_list = TopicPartitionList::from_rdkafka(partitions_ptr);
                Rebalance::Assign(topic_partition_list)
            },
            RDKafkaRespErr::RD_KAFKA_RESP_ERR__REVOKE_PARTITIONS => {
                Rebalance::Revoke
            },
            _ => {
                let error = unsafe { cstr_to_owned(rdkafka::rd_kafka_err2str(err)) };
                error!("Error rebalancing: {}", error);
                Rebalance::Error(error)
            }
        };

        self.pre_rebalance(&rebalance);

        // Execute rebalance
        unsafe {
            match err {
                RDKafkaRespErr::RD_KAFKA_RESP_ERR__ASSIGN_PARTITIONS => {
                    rdkafka::rd_kafka_assign(native_client.ptr(), partitions_ptr);
                },
                RDKafkaRespErr::RD_KAFKA_RESP_ERR__REVOKE_PARTITIONS => {
                    rdkafka::rd_kafka_assign(native_client.ptr(), ptr::null());
                },
                _ => {
                    rdkafka::rd_kafka_assign(native_client.ptr(), ptr::null());
                }
            }
        }
        self.post_rebalance(&rebalance);
    }

    /// Pre-rebalance callback. This method will run before the rebalance and should
    /// terminate its execution quickly.
    fn pre_rebalance(&self, _rebalance: &Rebalance) { }

    /// Post-rebalance callback. This method will run after the rebalance and should
    /// terminate its execution quickly.
    fn post_rebalance(&self, _rebalance: &Rebalance) { }
}

/// An empty consumer context that can be user when no context is needed.
#[derive(Clone)]
pub struct EmptyConsumerContext;

impl Context for EmptyConsumerContext { }
impl ConsumerContext for EmptyConsumerContext { }

/// Native rebalance callback. This callback will run on every rebalance, and it will call the
/// rebalance method defined in the current `Context`.
unsafe extern "C" fn rebalance_cb<C: ConsumerContext>(rk: *mut RDKafka,
                                                          err: RDKafkaRespErr,
                                                          partitions: *mut RDKafkaTopicPartitionList,
                                                          opaque_ptr: *mut c_void) {
    let context: &C = &*(opaque_ptr as *const C);
    let native_client = NativeClient::new(rk);

    context.rebalance(&native_client, err, partitions);
}

/// Specifies if the commit should be performed synchronously
/// or asynchronously.
pub enum CommitMode {
    /// Synchronous commit.
    Sync = 0,
    /// Asynchronous commit.
    Async = 1,
}

/// Common trait for all consumers.
pub trait Consumer<C: ConsumerContext> {
    /// Returns a reference to the BaseConsumer.
    fn get_base_consumer(&self) -> &BaseConsumer<C>;
    /// Returns a mutable reference to the BaseConsumer.
    fn get_base_consumer_mut(&mut self) -> &mut BaseConsumer<C>;

    // Default implementations

    /// Subscribe the consumer to a list of topics.
    fn subscribe(&mut self, topics: &Vec<&str>) -> KafkaResult<()> {
        self.get_base_consumer_mut().subscribe(topics)
    }

    /// Manually assign topics and partitions to the consumer.
    fn assign(&mut self, assignment: &TopicPartitionList) -> KafkaResult<()> {
        self.get_base_consumer_mut().assign(assignment)
    }

    /// Commit a specific message. If mode is set to CommitMode::Sync,
    /// the call will block until the message has been succesfully
    /// committed.
    fn commit_message(&self, message: &Message, mode: CommitMode) {
        self.get_base_consumer().commit_message(message, mode);
    }

    /// Returns the metadata information for all the topics in the cluster.
    fn fetch_metadata(&mut self, timeout_ms: i32) -> KafkaResult<Metadata> {
        self.get_base_consumer().fetch_metadata(timeout_ms)
    }
}
<commit_msg>Fix a whitespace mistake<commit_after>//! Base trait and common functionalities for all consumers.
extern crate rdkafka_sys as rdkafka;

pub mod base_consumer;
pub mod stream_consumer;

use self::rdkafka::types::*;

use std::ptr;
use std::os::raw::c_void;
use util::cstr_to_owned;

use client::{Context, NativeClient};
use message::Message;
use metadata::Metadata;
use error::KafkaResult;

pub use consumer::base_consumer::BaseConsumer;
pub use topic_partition_list::TopicPartitionList;

/// Rebalance information.
#[derive(Clone, Debug)]
pub enum Rebalance {
    Assign(TopicPartitionList),
    Revoke,
    Error(String)
}

/// Consumer specific Context. This user-defined object can be used to provide custom callbacks to
/// consumer events. Refer to the list of methods to check which callbacks can be specified.
pub trait ConsumerContext: Context {
    /// Implements the default rebalancing strategy and calls the `pre_rebalance` and
    /// `post_rebalance` methods. If this method is overridden, it will be responsibility
    /// of the user to call them if needed.
    fn rebalance(&self, native_client: &NativeClient, err: RDKafkaRespErr,
                 partitions_ptr: *mut RDKafkaTopicPartitionList) {

        let rebalance = match err {
            RDKafkaRespErr::RD_KAFKA_RESP_ERR__ASSIGN_PARTITIONS => {
                // TODO: this might be expensive
                let topic_partition_list = TopicPartitionList::from_rdkafka(partitions_ptr);
                Rebalance::Assign(topic_partition_list)
            },
            RDKafkaRespErr::RD_KAFKA_RESP_ERR__REVOKE_PARTITIONS => {
                Rebalance::Revoke
            },
            _ => {
                let error = unsafe { cstr_to_owned(rdkafka::rd_kafka_err2str(err)) };
                error!("Error rebalancing: {}", error);
                Rebalance::Error(error)
            }
        };

        self.pre_rebalance(&rebalance);

        // Execute rebalance
        unsafe {
            match err {
                RDKafkaRespErr::RD_KAFKA_RESP_ERR__ASSIGN_PARTITIONS => {
                    rdkafka::rd_kafka_assign(native_client.ptr(), partitions_ptr);
                },
                RDKafkaRespErr::RD_KAFKA_RESP_ERR__REVOKE_PARTITIONS => {
                    rdkafka::rd_kafka_assign(native_client.ptr(), ptr::null());
                },
                _ => {
                    rdkafka::rd_kafka_assign(native_client.ptr(), ptr::null());
                }
            }
        }
        self.post_rebalance(&rebalance);
    }

    /// Pre-rebalance callback. This method will run before the rebalance and should
    /// terminate its execution quickly.
    fn pre_rebalance(&self, _rebalance: &Rebalance) { }

    /// Post-rebalance callback. This method will run after the rebalance and should
    /// terminate its execution quickly.
    fn post_rebalance(&self, _rebalance: &Rebalance) { }
}

/// An empty consumer context that can be user when no context is needed.
#[derive(Clone)]
pub struct EmptyConsumerContext;

impl Context for EmptyConsumerContext { }
impl ConsumerContext for EmptyConsumerContext { }

/// Native rebalance callback. This callback will run on every rebalance, and it will call the
/// rebalance method defined in the current `Context`.
unsafe extern "C" fn rebalance_cb<C: ConsumerContext>(rk: *mut RDKafka,
                                                      err: RDKafkaRespErr,
                                                      partitions: *mut RDKafkaTopicPartitionList,
                                                      opaque_ptr: *mut c_void) {
    let context: &C = &*(opaque_ptr as *const C);
    let native_client = NativeClient::new(rk);

    context.rebalance(&native_client, err, partitions);
}

/// Specifies if the commit should be performed synchronously
/// or asynchronously.
pub enum CommitMode {
    /// Synchronous commit.
    Sync = 0,
    /// Asynchronous commit.
    Async = 1,
}

/// Common trait for all consumers.
pub trait Consumer<C: ConsumerContext> {
    /// Returns a reference to the BaseConsumer.
    fn get_base_consumer(&self) -> &BaseConsumer<C>;
    /// Returns a mutable reference to the BaseConsumer.
    fn get_base_consumer_mut(&mut self) -> &mut BaseConsumer<C>;

    // Default implementations

    /// Subscribe the consumer to a list of topics.
    fn subscribe(&mut self, topics: &Vec<&str>) -> KafkaResult<()> {
        self.get_base_consumer_mut().subscribe(topics)
    }

    /// Manually assign topics and partitions to the consumer.
    fn assign(&mut self, assignment: &TopicPartitionList) -> KafkaResult<()> {
        self.get_base_consumer_mut().assign(assignment)
    }

    /// Commit a specific message. If mode is set to CommitMode::Sync,
    /// the call will block until the message has been succesfully
    /// committed.
    fn commit_message(&self, message: &Message, mode: CommitMode) {
        self.get_base_consumer().commit_message(message, mode);
    }

    /// Returns the metadata information for all the topics in the cluster.
    fn fetch_metadata(&mut self, timeout_ms: i32) -> KafkaResult<Metadata> {
        self.get_base_consumer().fetch_metadata(timeout_ms)
    }
}
<|endoftext|>


<commit_before>/// Represents a context tree.
pub struct ContextTree {
  root: Node,
}

impl ContextTree {
  pub fn create(depth: usize) -> ContextTree {
    ContextTree {
      root: Node::new(depth, 0),
    }
  }

  /// Returns the total number of nodes in the tree.
  pub fn size(&self) -> usize {
    self.root.size()
  }
}

/// One node in the context tree.
struct Node {
  left: Option<Box<Node>>,
  right: Option<Box<Node>>,
}

impl Node {
  fn new(depth: usize, level: usize) -> Node {
    return if level == depth {
      return Node::new_leaf();
    } else {
      return Node::new_inner(depth, level + 1);
    }
  }

  fn new_leaf() -> Node {
    Node {
      left: None,
      right: None,
    }
  }

  fn new_inner(depth: usize, level: usize) -> Node {
    Node {
      left: Some(Box::new(Node::new(depth, level + 1))),
      right: Some(Box::new(Node::new(depth, level + 1))),
    }
  }

  /// Returns the number of nodes in the subtree rooted at this node.
  fn size(&self) -> usize {
    return if self.is_leaf() { 1 } else {
      let left_size = self.left.as_ref().unwrap().size();
      let right_size = self.right.as_ref().unwrap().size();
      left_size + right_size + 1
    };
  }

  fn is_leaf(&self) -> bool {
    if self.left.is_some() != self.right.is_some() {
      panic!("ContextTree nodes cannot have only one child.");
    }
    return !self.left.is_some();
  }
}
<commit_msg>Rename context tree children to make it clear which is which.<commit_after>/// Represents a context tree.
pub struct ContextTree {
  root: Node,
}

impl ContextTree {
  pub fn create(depth: usize) -> ContextTree {
    ContextTree {
      root: Node::new(depth, 0),
    }
  }

  /// Returns the total number of nodes in the tree.
  pub fn size(&self) -> usize {
    self.root.size()
  }
}

/// One node in the context tree.
struct Node {
  zero_child: Option<Box<Node>>,
  one_child: Option<Box<Node>>,
}

impl Node {
  fn new(depth: usize, level: usize) -> Node {
    return if level == depth {
      return Node::new_leaf();
    } else {
      return Node::new_inner(depth, level + 1);
    }
  }

  fn new_leaf() -> Node {
    Node {
      zero_child: None,
      one_child: None,
    }
  }

  fn new_inner(depth: usize, level: usize) -> Node {
    Node {
      zero_child: Some(Box::new(Node::new(depth, level + 1))),
      one_child: Some(Box::new(Node::new(depth, level + 1))),
    }
  }

  /// Returns the number of nodes in the subtree rooted at this node.
  fn size(&self) -> usize {
    return if self.is_leaf() { 1 } else {
      let zero_child_size = self.zero_child.as_ref().unwrap().size();
      let one_child_size = self.one_child.as_ref().unwrap().size();
      zero_child_size + one_child_size + 1
    };
  }

  fn is_leaf(&self) -> bool {
    if self.zero_child.is_some() != self.one_child.is_some() {
      panic!("ContextTree nodes cannot have only one child.");
    }
    return !self.zero_child.is_some();
  }
}
<|endoftext|>


<commit_before>use serde;
use serde_json;
use std::io;
use std::string;

error_chain! {
    types {
        Error, ErrorKind, Result;
    }

    links {
    }

    foreign_links {
        io::Error, IO;
        string::FromUtf8Error, Utf8;
        serde_json::Error, JsonDecode;
    }

    errors {
        NoSuchType(type_name: String) {
            description("no such type")
            display("no such type: {:?}", type_name)
        }
        InvalidSchema {
            description("invalid schema")
            display("invalid schema")
        }
        BadFileMagic(magic: Vec<u8>) {
            description("bad file magic")
            display("bad file magic: {:?}", magic)
        }
        NoSchema {
            description("the container has no schema")
            display("the container has no schema")
        }
        NoRootType {
            description("the container schema has no root type")
            display("the container schema has no root type")
        }
        UnsupportedCodec(codec: String) {
            description("unsupported codec")
            display("unsupported codec: {:?}", codec)
        }
        FieldTypeMismatch(field: &'static str, expected: &'static str) {
            description("unexpected field type")
            display("the value for {:?} was not a {}", field, expected)
        }
        RequiredFieldMissing(field: &'static str) {
            description("required field missing")
            display("the {:?} field is required", field)
        }
        DuplicateSchema(name: String) {
            description("there are two schemata with the same name")
            display("there are two schemata called {:?}", name)
        }
        EndOfStream {
            description("end of stream")
            display("end of stream")
        }
        IntegerOverflow {
            description("integer overflow")
            display("integer overflow")
        }
        NegativeLength {
            description("negative length")
            display("negative length")
        }
    }
}

impl serde::Error for Error {
    fn custom<S>(msg: S) -> Error
        where S: Into<String>
    {
        msg.into().into()
    }

    fn end_of_stream() -> Error {
        ErrorKind::EndOfStream.into()
    }
}
<commit_msg>Remove unnecessary error-chain sections<commit_after>use serde;
use serde_json;
use std::io;
use std::string;

error_chain! {
    foreign_links {
        io::Error, IO;
        string::FromUtf8Error, Utf8;
        serde_json::Error, JsonDecode;
    }

    errors {
        NoSuchType(type_name: String) {
            description("no such type")
            display("no such type: {:?}", type_name)
        }
        InvalidSchema {
            description("invalid schema")
            display("invalid schema")
        }
        BadFileMagic(magic: Vec<u8>) {
            description("bad file magic")
            display("bad file magic: {:?}", magic)
        }
        NoSchema {
            description("the container has no schema")
            display("the container has no schema")
        }
        NoRootType {
            description("the container schema has no root type")
            display("the container schema has no root type")
        }
        UnsupportedCodec(codec: String) {
            description("unsupported codec")
            display("unsupported codec: {:?}", codec)
        }
        FieldTypeMismatch(field: &'static str, expected: &'static str) {
            description("unexpected field type")
            display("the value for {:?} was not a {}", field, expected)
        }
        RequiredFieldMissing(field: &'static str) {
            description("required field missing")
            display("the {:?} field is required", field)
        }
        DuplicateSchema(name: String) {
            description("there are two schemata with the same name")
            display("there are two schemata called {:?}", name)
        }
        EndOfStream {
            description("end of stream")
            display("end of stream")
        }
        IntegerOverflow {
            description("integer overflow")
            display("integer overflow")
        }
        NegativeLength {
            description("negative length")
            display("negative length")
        }
    }
}

impl serde::Error for Error {
    fn custom<S>(msg: S) -> Error
        where S: Into<String>
    {
        msg.into().into()
    }

    fn end_of_stream() -> Error {
        ErrorKind::EndOfStream.into()
    }
}
<|endoftext|>


<commit_before>extern crate ocl;
#[allow(non_camel_case_types)]
#[allow(non_upper_case_globals)]
#[allow(non_snake_case)]
pub mod ffi;
use std::mem;
pub use ffi::clfftSetupData;

macro_rules! clfft_try {
    ( $result_expr: expr) => {
        let result = $result_expr;
        if result != ffi::clfftStatus::CLFFT_SUCCESS {
            let error_code: i32 = unsafe { mem::transmute(result) };
            return ocl::Error::err(format!("FOREIGN_CLFFT_ERROR {}\nSee http://clmathlibraries.github.io/clFFT/clFFT_8h.html#a74e303bed132064cfa22c1cce96d2bce for further information", error_code));
        }
    }
}

macro_rules! clfft_panic {
    ( $result_expr: expr) => {
        let result = $result_expr;
        if result != ffi::clfftStatus::CLFFT_SUCCESS {
            panic!("Unexpeced error in foreign library: {:?}", result);
        }
    }
}

/// Initialize the internal FFT resources such as FFT implementation caches kernels, programs, and buffers.
pub fn init_lib() -> ffi::clfftSetupData {
    let mut major: ocl::ffi::cl_uint = 0;
    let mut minor: ocl::ffi::cl_uint = 0;
    let mut patch: ocl::ffi::cl_uint = 0;
    clfft_panic!( unsafe { ffi::clfftGetVersion(&mut major, &mut minor, &mut patch) } );
    let data = ffi::clfftSetupData {
        major: major,
        minor: minor,
        patch: patch,
        debugFlags: 0,
    };

    clfft_panic!(unsafe { ffi::clfftSetup(&data) });
    data
}

/// Release all internal resources acquired during `init_lib`.
pub fn drop_lib() {
    unsafe { ffi::clfftTeardown() };
}

fn translate_to_fft_dim(dims: ocl::SpatialDims) -> ffi::clfftDim {
    match dims.dim_count() {
        1 => ffi::clfftDim::CLFFT_1D,
        2 => ffi::clfftDim::CLFFT_2D,
        3 => ffi::clfftDim::CLFFT_3D,
        n => panic!("Number of dimensions must be 1, 2 or 3, but it is {}", n)
    }
}

pub struct FftPlan {
    handle: ffi::clfftPlanHandle
}

/// Specify the expected precision of each FFT.
pub enum Precision {
    Single,
    Double,
    SingleFast,
    DoubleFast
}

fn translate_precision(precision: Precision) -> ffi::clfftPrecision {
    match precision {
        Precision::Single => ffi::clfftPrecision::CLFFT_SINGLE,
        Precision::Double => ffi::clfftPrecision::CLFFT_DOUBLE,
        Precision::SingleFast => ffi::clfftPrecision::CLFFT_SINGLE_FAST,
        Precision::DoubleFast => ffi::clfftPrecision::CLFFT_DOUBLE_FAST
    }
}

fn translate_precision_back(precision: ffi::clfftPrecision) -> Precision {
    match precision {
        ffi::clfftPrecision::CLFFT_SINGLE => Precision::Single,
        ffi::clfftPrecision::CLFFT_DOUBLE => Precision::Double,
        ffi::clfftPrecision::CLFFT_SINGLE_FAST => Precision::SingleFast,
        ffi::clfftPrecision::CLFFT_DOUBLE_FAST => Precision::DoubleFast,
        ffi::clfftPrecision::ENDPRECISION => panic!("ENDPRECISION should never be returned")
    }
}

impl FftPlan {
    /// Create a plan object initialized entirely with default values.
    ///
    /// A plan is a repository of state for calculating FFT's.  Allows the runtime to pre-calculate kernels, programs
    /// and buffers and associate them with buffers of specified dimensions.
    pub fn default<D: Into<ocl::SpatialDims>>(pro_que: &ocl::ProQue, dims: D)
        -> ocl::Result<FftPlan> {
        let context = unsafe { pro_que.context().core_as_ref().as_ptr() };
        let dims = dims.into();
        let dim = translate_to_fft_dim(dims);
        let lengths = try!(dims.to_lens());
        let mut plan: ffi::clfftPlanHandle = 0;
        clfft_try!( unsafe { ffi::clfftCreateDefaultPlan(&mut plan, context, dim, &lengths as *const usize) } );
        Ok(FftPlan {  handle: plan })
    }

    /// Returns the native clFFT plan handle.
    pub fn plan_handle(&self) -> ffi::clfftPlanHandle {
        self.handle
    }

    /// Set the floating point precision of the FFT data
    pub fn set_precision(&mut self, precision: Precision) -> ocl::Result<()> {
        let precision = translate_precision(precision);
        clfft_try!(unsafe { ffi::clfftSetPlanPrecision(self.handle, precision) });
        Ok(())
    }

    pub fn get_precision(&self) -> ocl::Result<Precision> {
        let mut precision = ffi::clfftPrecision::CLFFT_SINGLE;
        clfft_try!(unsafe { ffi::clfftGetPlanPrecision(self.handle, &mut precision) });
        Ok(translate_precision_back(precision))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ocl::{ProQue};
    use ocl::builders::ProgramBuilder;

    #[test]
    pub fn fft_test() {
        init_lib();
        let prog_bldr = ProgramBuilder::new();
        let ocl_pq = ProQue::builder()
            .prog_bldr(prog_bldr)
            .build()
            .expect("Building ProQue");
        FftPlan::default(&ocl_pq, [1000]);
        drop_lib();
    }
}<commit_msg>More options for plasn<commit_after>extern crate ocl;
#[allow(non_camel_case_types)]
#[allow(non_upper_case_globals)]
#[allow(non_snake_case)]
pub mod ffi;
use std::mem;
pub use ffi::clfftSetupData;

macro_rules! clfft_try {
    ( $result_expr: expr) => {
        let result = $result_expr;
        if result != ffi::clfftStatus::CLFFT_SUCCESS {
            let error_code: i32 = unsafe { mem::transmute(result) };
            return ocl::Error::err(format!("FOREIGN_CLFFT_ERROR {}\nSee http://clmathlibraries.github.io/clFFT/clFFT_8h.html#a74e303bed132064cfa22c1cce96d2bce for further information", error_code));
        }
    }
}

macro_rules! clfft_panic {
    ( $result_expr: expr) => {
        let result = $result_expr;
        if result != ffi::clfftStatus::CLFFT_SUCCESS {
            panic!("Unexpeced error in foreign library: {:?}", result);
        }
    }
}

/// Initialize the internal FFT resources such as FFT implementation caches kernels, programs, and buffers.
pub fn init_lib() -> ffi::clfftSetupData {
    let mut major: ocl::ffi::cl_uint = 0;
    let mut minor: ocl::ffi::cl_uint = 0;
    let mut patch: ocl::ffi::cl_uint = 0;
    clfft_panic!( unsafe { ffi::clfftGetVersion(&mut major, &mut minor, &mut patch) } );
    let data = ffi::clfftSetupData {
        major: major,
        minor: minor,
        patch: patch,
        debugFlags: 0,
    };

    clfft_panic!(unsafe { ffi::clfftSetup(&data) });
    data
}

/// Release all internal resources acquired during `init_lib`.
pub fn drop_lib() {
    unsafe { ffi::clfftTeardown() };
}

fn translate_to_fft_dim(dims: ocl::SpatialDims) -> ffi::clfftDim {
    match dims.dim_count() {
        1 => ffi::clfftDim::CLFFT_1D,
        2 => ffi::clfftDim::CLFFT_2D,
        3 => ffi::clfftDim::CLFFT_3D,
        n => panic!("Number of dimensions must be 1, 2 or 3, but it is {}", n)
    }
}

pub struct FftPlan {
    handle: ffi::clfftPlanHandle
}

/// Specify the expected precision of each FFT.
pub enum Precision {
    Single,
    Double,
    SingleFast,
    DoubleFast
}

fn translate_precision(precision: Precision) -> ffi::clfftPrecision {
    match precision {
        Precision::Single => ffi::clfftPrecision::CLFFT_SINGLE,
        Precision::Double => ffi::clfftPrecision::CLFFT_DOUBLE,
        Precision::SingleFast => ffi::clfftPrecision::CLFFT_SINGLE_FAST,
        Precision::DoubleFast => ffi::clfftPrecision::CLFFT_DOUBLE_FAST
    }
}

fn translate_precision_back(precision: ffi::clfftPrecision) -> Precision {
    match precision {
        ffi::clfftPrecision::CLFFT_SINGLE => Precision::Single,
        ffi::clfftPrecision::CLFFT_DOUBLE => Precision::Double,
        ffi::clfftPrecision::CLFFT_SINGLE_FAST => Precision::SingleFast,
        ffi::clfftPrecision::CLFFT_DOUBLE_FAST => Precision::DoubleFast,
        ffi::clfftPrecision::ENDPRECISION => panic!("ENDPRECISION should never be returned")
    }
}

/// Specify the expected layouts of the buffers.
pub enum Layout {
    ComplexInterleaved,
    ComplexPlanar,
    HermitianInterleaved,
    HermitianPlanar,
    Real
}

fn translate_layout(layout: Layout) -> ffi::clfftLayout {
    match layout {
        Layout::ComplexInterleaved => ffi::clfftLayout::CLFFT_COMPLEX_INTERLEAVED,
        Layout::ComplexPlanar => ffi::clfftLayout::CLFFT_COMPLEX_PLANAR,
        Layout::HermitianInterleaved => ffi::clfftLayout::CLFFT_HERMITIAN_INTERLEAVED,
        Layout::HermitianPlanar => ffi::clfftLayout::CLFFT_HERMITIAN_PLANAR,
        Layout::Real => ffi::clfftLayout::CLFFT_REAL
    }
}

fn translate_layout_back(layout: ffi::clfftLayout) -> Layout {
    match layout {
        ffi::clfftLayout::CLFFT_COMPLEX_INTERLEAVED => Layout::ComplexInterleaved,
        ffi::clfftLayout::CLFFT_COMPLEX_PLANAR => Layout::ComplexPlanar,
        ffi::clfftLayout::CLFFT_HERMITIAN_INTERLEAVED => Layout::HermitianInterleaved,
        ffi::clfftLayout::CLFFT_HERMITIAN_PLANAR => Layout::HermitianPlanar,
        ffi::clfftLayout::CLFFT_REAL => Layout::Real,
        ffi::clfftLayout::ENDLAYOUT => panic!("ENDLAYOUT should never be returned")
    }
}

/// Specify the expected direction of each FFT, time or the frequency domains
pub enum Direction {
    Forward,
    Backward
}

fn translate_direction(direction: Direction) -> ffi::clfftDirection {
    match direction {
        Direction::Forward => ffi::clfftDirection::CLFFT_FORWARD,
        Direction::Backward => ffi::clfftDirection::CLFFT_BACKWARD
    }
}

fn translate_direction_back(direction: ffi::clfftDirection) -> Direction {
    match direction {
        ffi::clfftDirection::CLFFT_FORWARD => Direction::Forward,
        ffi::clfftDirection::CLFFT_BACKWARD => Direction::Backward,
        ffi::clfftDirection::ENDDIRECTION => panic!("ENDDIRECTION should never be returned")
    }
}

/// pecify wheter the input buffers are overwritten with results
pub enum Location {
    Inplace,
    OutOfPlace
}

fn translate_location(location: Location) -> ffi::clfftResultLocation {
    match location {
        Location::Inplace => ffi::clfftResultLocation::CLFFT_INPLACE,
        Location::OutOfPlace => ffi::clfftResultLocation::CLFFT_OUTOFPLACE
    }
}

fn translate_location_back(location: ffi::clfftResultLocation) -> Location {
    match location {
        ffi::clfftResultLocation::CLFFT_INPLACE => Location::Inplace,
        ffi::clfftResultLocation::CLFFT_OUTOFPLACE => Location::OutOfPlace,
        ffi::clfftResultLocation::ENDPLACE => panic!("ENDPLACE should never be returned")
    }
}

impl FftPlan {
    /// Create a plan object initialized entirely with default values.
    ///
    /// A plan is a repository of state for calculating FFT's.  Allows the runtime to pre-calculate kernels, programs
    /// and buffers and associate them with buffers of specified dimensions.
    pub fn default<D: Into<ocl::SpatialDims>>(pro_que: &ocl::ProQue, dims: D)
        -> ocl::Result<FftPlan> {
        let context = unsafe { pro_que.context().core_as_ref().as_ptr() };
        let dims = dims.into();
        let dim = translate_to_fft_dim(dims);
        let lengths = try!(dims.to_lens());
        let mut plan: ffi::clfftPlanHandle = 0;
        clfft_try!( unsafe { ffi::clfftCreateDefaultPlan(&mut plan, context, dim, &lengths as *const usize) } );
        Ok(FftPlan {  handle: plan })
    }

    /// Returns the native clFFT plan handle.
    pub fn plan_handle(&self) -> ffi::clfftPlanHandle {
        self.handle
    }

    /// Set the floating point precision of the FFT data
    pub fn set_precision(&mut self, precision: Precision) -> ocl::Result<()> {
        let precision = translate_precision(precision);
        clfft_try!(unsafe { ffi::clfftSetPlanPrecision(self.handle, precision) });
        Ok(())
    }

    pub fn get_precision(&self) -> ocl::Result<Precision> {
        let mut precision = ffi::clfftPrecision::CLFFT_SINGLE;
        clfft_try!(unsafe { ffi::clfftGetPlanPrecision(self.handle, &mut precision) });
        Ok(translate_precision_back(precision))
    }

    /// Set the expected layout of the input and output buffers
    pub fn set_layout(&mut self, input_layout: Layout, output_layout: Layout) -> ocl::Result<()> {
        let input_layout = translate_layout(input_layout);
        let output_layout = translate_layout(output_layout);
        clfft_try!(unsafe { ffi::clfftSetLayout(self.handle, input_layout, output_layout) });
        Ok(())
    }

    pub fn get_layout(&self) -> ocl::Result<(Layout, Layout)> {
        let mut input_layout = ffi::clfftLayout::CLFFT_COMPLEX_INTERLEAVED;
        let mut output_layout = ffi::clfftLayout::CLFFT_COMPLEX_INTERLEAVED;
        clfft_try!(unsafe { ffi::clfftGetLayout(self.handle, &mut input_layout, &mut output_layout) });
        Ok((translate_layout_back(input_layout), translate_layout_back(output_layout)))
    }

    /// Set whether the input buffers are to be overwritten with results
    pub fn set_result_location(&mut self, location: Location) -> ocl::Result<()> {
        let location = translate_location(location);
        clfft_try!(unsafe { ffi::clfftSetResultLocation(self.handle, location) });
        Ok(())
    }

    pub fn get_result_location(&self) -> ocl::Result<Location> {
        let mut location = ffi::clfftResultLocation::CLFFT_INPLACE;
        clfft_try!(unsafe { ffi::clfftGetResultLocation(self.handle, &mut location) });
        Ok(translate_location_back(location))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ocl::{ProQue};
    use ocl::builders::ProgramBuilder;

    #[test]
    pub fn fft_test() {
        init_lib();
        let prog_bldr = ProgramBuilder::new();
        let ocl_pq = ProQue::builder()
            .prog_bldr(prog_bldr)
            .build()
            .expect("Building ProQue");
        FftPlan::default(&ocl_pq, [1000]);
        drop_lib();
    }
}<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::thread;
use std::time::{Duration, Instant};

use anyhow::Result;
use clap::Clap;
use laminar::{Packet, Socket};
use legion::*;
use log::{debug, info, LevelFilter};
use macroquad::prelude::*;
use simplelog::{Config, TermLogger, TerminalMode};

use zombenum_shared::components::GameEntityId;
use zombenum_shared::components::{Component, PlayerId};
use zombenum_shared::messages::client::ClientMessage;
use zombenum_shared::messages::server::{ServerMessage, ServerMessageType};

mod cli;
mod input;
mod network;
mod systems;

use input::get_input_events;
use network::receive::get_server_updates;
use network::send::send_input_events;

const SERVER: &str = "127.0.0.1:14191";

fn window_conf() -> Conf {
    Conf {
        window_title: "Zombenum".to_string(),
        ..Default::default()
    }
}

#[macroquad::main(window_conf)]
async fn main() {
    run().await.expect("Error in run()");
}

async fn run() -> Result<()> {
    let opt = cli::CliArguments::parse();
    // Set the verbosity level of the logger.
    let level = match opt.verbose {
        0 => LevelFilter::Error,
        1 => LevelFilter::Warn,
        2 => LevelFilter::Info,
        _ => LevelFilter::Debug,
    };
    TermLogger::init(level, Config::default(), TerminalMode::Mixed)?;

    info!("Starting client");

    let mut socket = Socket::bind_any()?;
    info!("Bound on {}", socket.local_addr()?.to_string());

    let (sender, receiver) = (socket.get_packet_sender(), socket.get_event_receiver());
    let _thread = thread::spawn(move || socket.start_polling());

    info!("Server is {}", SERVER);
    let server = SERVER.parse()?;

    let mut world = World::default();
    let mut resources = Resources::default();
    let mut schedule = Schedule::builder()
        .add_system(systems::draw_system())
        .build();

    // Legion's ids of the same entities, might differ between the server and client worlds.
    // This map allows us to easily map GameEntityIds, which are shared and identical between all
    // members, to our local Legion Entity ids. This is, for instance, used during Component updates.
    let mut game_entity_id_map: HashMap<GameEntityId, Entity> = HashMap::new();

    // We start out not connected. Eventually, we'll send a `JoinGame` message and receive a reply
    // with our first initial game state. That's when we'll first be connected. However, we might
    // run into issues and the server might time us out in which case we'll receive network errors
    // and we'll revert this value to `false` and try to connect again.
    // let mut is_connected = false;

    let mut last_server_session_time = Duration::default();
    let mut last_server_session_time_changed_at = Instant::now();

    info!("Sending JoinGame");
    let join_game_msg = ClientMessage::JoinGame(PlayerId::new());
    sender.send(Packet::reliable_unordered(
        server,
        bincode::serialize(&join_game_msg)?,
    ))?;

    info!("Start main loop");
    loop {
        let session_time =
            last_server_session_time + (Instant::now() - last_server_session_time_changed_at);
        debug!("Starting tick at session time {:?}", session_time);
        // Get a list of all input events issued by the user.
        let input_events = get_input_events();

        // Send those input events to the server.
        send_input_events(&input_events, &sender, server, session_time)?;

        // Collect and handle all server messages since the last tick.
        let server_messages = get_server_updates(&receiver, server)?;
        for message in server_messages {
            let received_server_session_time =
                handle_server_update(&mut world, &mut game_entity_id_map, message);

            // Update last received server session time if it is more recent than the current value.
            if last_server_session_time < received_server_session_time {
                last_server_session_time = received_server_session_time;
                last_server_session_time_changed_at = Instant::now();
                debug!(
                    "Received new server session time {} at {:?}",
                    last_server_session_time.as_millis(),
                    last_server_session_time_changed_at,
                );
            }
        }

        clear_background(WHITE);
        schedule.execute(&mut world, &mut resources);
        next_frame().await
    }
}

fn handle_server_update(
    world: &mut World,
    game_entity_id_map: &mut HashMap<GameEntityId, Entity>,
    message: ServerMessage,
) -> Duration {
    if let ServerMessageType::WorldUpdate(world_update) = message.message_type {
        for components in world_update.new_entities {
            let entity = world.push(());
            let mut entry = world.entry(entity).unwrap();

            // Get the GameEntityId from the entity's component list.
            // We have to go through all components and filter the correct one.
            let game_entity_id = components
                .iter()
                .filter(|component| matches!(component, Component::GameEntityId(_)))
                .map(|component| {
                    if let Component::GameEntityId(component) = component {
                        Some(component.clone())
                    } else {
                        None
                    }
                })
                .next();

            // Add the GameEntityId with our local EntityId to the game_entity_id_map.
            // Print an error message if there's no GameEntityId.
            // This should never be the case!
            if let Some(game_entity_id) = game_entity_id.flatten() {
                game_entity_id_map.insert(game_entity_id, entity);
            } else {
                error!("Got an entity without an GameEntityId: {:?}", components);
            }

            info!("Create new entity");
            for component in components {
                info!("Add component: {:?}", component);
                add_component_to_entity(&mut entry, component);
            }
        }
    }
    message.session_time
}

fn add_component_to_entity(entry: &mut world::Entry, component: Component) {
    match component {
        Component::GameEntityId(component) => entry.add_component(component),
        Component::PlayerId(component) => entry.add_component(component),
        Component::Position(component) => entry.add_component(component),
        Component::Drawable(component) => entry.add_component(component),
    }
}
<commit_msg>Update components<commit_after>use std::collections::HashMap;
use std::thread;
use std::time::{Duration, Instant};

use anyhow::Result;
use clap::Clap;
use laminar::{Packet, Socket};
use legion::*;
use log::{debug, error, info, warn, LevelFilter};
use macroquad::prelude::*;
use simplelog::{Config, TermLogger, TerminalMode};

use zombenum_shared::components::GameEntityId;
use zombenum_shared::components::{Component, PlayerId};
use zombenum_shared::messages::client::ClientMessage;
use zombenum_shared::messages::server::{ServerMessage, ServerMessageType, WorldUpdateMessage};

mod cli;
mod input;
mod network;
mod systems;

use input::get_input_events;
use network::receive::get_server_updates;
use network::send::send_input_events;

const SERVER: &str = "127.0.0.1:14191";

fn window_conf() -> Conf {
    Conf {
        window_title: "Zombenum".to_string(),
        ..Default::default()
    }
}

#[macroquad::main(window_conf)]
async fn main() {
    run().await.expect("Error in run()");
}

async fn run() -> Result<()> {
    let opt = cli::CliArguments::parse();
    // Set the verbosity level of the logger.
    let level = match opt.verbose {
        0 => LevelFilter::Error,
        1 => LevelFilter::Warn,
        2 => LevelFilter::Info,
        _ => LevelFilter::Debug,
    };
    TermLogger::init(level, Config::default(), TerminalMode::Mixed)?;

    info!("Starting client");

    let mut socket = Socket::bind_any()?;
    info!("Bound on {}", socket.local_addr()?.to_string());

    let (sender, receiver) = (socket.get_packet_sender(), socket.get_event_receiver());
    let _thread = thread::spawn(move || socket.start_polling());

    info!("Server is {}", SERVER);
    let server = SERVER.parse()?;

    let mut world = World::default();
    let mut resources = Resources::default();
    let mut schedule = Schedule::builder()
        .add_system(systems::draw_system())
        .build();

    // Legion's ids of the same entities, might differ between the server and client worlds.
    // This map allows us to easily map GameEntityIds, which are shared and identical between all
    // members, to our local Legion Entity ids. This is, for instance, used during Component updates.
    let mut game_entity_id_map: HashMap<GameEntityId, Entity> = HashMap::new();

    // We start out not connected. Eventually, we'll send a `JoinGame` message and receive a reply
    // with our first initial game state. That's when we'll first be connected. However, we might
    // run into issues and the server might time us out in which case we'll receive network errors
    // and we'll revert this value to `false` and try to connect again.
    // let mut is_connected = false;

    let mut last_server_session_time = Duration::default();
    let mut last_server_session_time_changed_at = Instant::now();

    info!("Sending JoinGame");
    let join_game_msg = ClientMessage::JoinGame(PlayerId::new());
    sender.send(Packet::reliable_unordered(
        server,
        bincode::serialize(&join_game_msg)?,
    ))?;

    info!("Start main loop");
    loop {
        let session_time =
            last_server_session_time + (Instant::now() - last_server_session_time_changed_at);
        debug!("Starting tick at session time {:?}", session_time);
        // Get a list of all input events issued by the user.
        let input_events = get_input_events();

        // Send those input events to the server.
        send_input_events(&input_events, &sender, server, session_time)?;

        // Collect and handle all server messages since the last tick.
        let server_messages = get_server_updates(&receiver, server)?;
        for message in server_messages {
            let received_server_session_time =
                handle_world_update(&mut world, &mut game_entity_id_map, message);

            // Update last received server session time if it is more recent than the current value.
            if last_server_session_time < received_server_session_time {
                last_server_session_time = received_server_session_time;
                last_server_session_time_changed_at = Instant::now();
                debug!(
                    "Received new server session time {} at {:?}",
                    last_server_session_time.as_millis(),
                    last_server_session_time_changed_at,
                );
            }
        }

        clear_background(WHITE);
        schedule.execute(&mut world, &mut resources);
        next_frame().await
    }
}

fn handle_world_update(
    world: &mut World,
    game_entity_id_map: &mut HashMap<GameEntityId, Entity>,
    message: ServerMessage,
) -> Duration {
    if let ServerMessageType::WorldUpdate(world_update) = message.message_type {
        add_new_entities(world, game_entity_id_map, world_update.new_entities);
        update_components(world, game_entity_id_map, world_update.changed_components);
    }

    message.session_time
}

/// Some component changed on the server.
/// This function handles those updates and integrates them into our world.
fn update_components(
    world: &mut World,
    game_entity_id_map: &mut HashMap<GameEntityId, Entity>,
    changed_components: Vec<(GameEntityId, Component)>,
) {
    for (game_entity_id, component) in changed_components {
        let entity = if let Some(entity) = game_entity_id_map.get(&game_entity_id) {
            entity.clone()
        } else {
            warn!(
                "Got update {:?} for unknown entity {:?}",
                component, game_entity_id
            );
            continue;
        };

        // For now this logic is SUPER simple, this will be replaced in the future
        // By adding a component over an existing one, we just delete the old component and replace
        // it with the new values wie got by the server.
        if let Some(mut entry) = world.entry(entity) {
            add_component_to_entity(&mut entry, component);
        } else {
            error!(
                "Got update {:?} for unknown entity {:?}",
                component, game_entity_id
            );
        }
    }
}

/// The server added some new entities.
/// This functions adds all those entities to our world and registers them in our
/// game_entity_id_map.
fn add_new_entities(
    world: &mut World,
    game_entity_id_map: &mut HashMap<GameEntityId, Entity>,
    entities: Vec<Vec<Component>>,
) {
    for components in entities.into_iter() {
        let entity = world.push(());
        let mut entry = world.entry(entity).unwrap();

        // Get the GameEntityId from the entity's component list.
        // We have to go through all components and filter the correct one.
        let game_entity_id = components
            .iter()
            .filter(|component| matches!(component, Component::GameEntityId(_)))
            .map(|component| {
                if let Component::GameEntityId(component) = component {
                    Some(component.clone())
                } else {
                    None
                }
            })
            .next();

        // Add the GameEntityId with our local EntityId to the game_entity_id_map.
        // Print an error message if there's no GameEntityId.
        // This should never be the case!
        if let Some(game_entity_id) = game_entity_id.flatten() {
            game_entity_id_map.insert(game_entity_id, entity);
        } else {
            error!("Got an entity without an GameEntityId: {:?}", components);
        }

        info!("Create new entity");
        for component in components {
            info!("Add component: {:?}", component);
            add_component_to_entity(&mut entry, component);
        }
    }
}

fn add_component_to_entity(entry: &mut world::Entry, component: Component) {
    match component {
        Component::GameEntityId(component) => entry.add_component(component),
        Component::PlayerId(component) => entry.add_component(component),
        Component::Position(component) => entry.add_component(component),
        Component::Drawable(component) => entry.add_component(component),
    }
}
<|endoftext|>