<commit_before>use cgmath::*;
use crossbeam;
use rand;
use rand::Rng;
use utils::v3;

pub type V3 = Vector3<f32>;

pub struct CubeState {
    pub pos: V3,
    pub rot: Basis3<f32>,
    pub spd: V3,
    pub rot_axe: V3,
    pub rot_spd: f32,
}

pub struct State {
    pub tick: f64,
    pub cubes: Vec<CubeState>,
    pub light_pos: V3,
}

impl State {
    pub fn new(object_count: u32) -> Self {
        let mut rng = rand::XorShiftRng::new_unseeded();
        let mut cubes = Vec::with_capacity(object_count as usize);
        for _ in 0 .. object_count {
            let cube_state = CubeState {
                pos: v3((rng.next_f32() - 0.5) * 200.,
                             (rng.next_f32() - 0.5) * 200.,
                             (rng.next_f32() - 1.) * 200.),
                rot: <Basis3<f32> as Rotation<_>>::one(),
                spd: v3(rng.next_f32() - 0.5,
                             rng.next_f32() - 0.5,
                             rng.next_f32() - 0.5).normalize()*rng.next_f32(),
                rot_axe: v3(rng.next_f32() - 0.5,
                             rng.next_f32() - 0.5,
                             rng.next_f32() - 0.5)
                              .normalize(),
                rot_spd: rng.next_f32()*90.0,
            };
            cubes.push(cube_state);
        };
        State {
            tick: 0.0,
            cubes: cubes,
            light_pos: v3(0., 0., 0.),
        }
    }

    pub fn update(&self, dt: f32, thread_cnt: u32) -> Self {
        // advance position of cubes
        let mut cubes = Vec::with_capacity(self.cubes.len());
        // Uninitialized content of cubes will be immediately overwritten
        unsafe{ cubes.set_len(self.cubes.len()) };
        crossbeam::scope(|scope| {
            let chunk_len = (cubes.len()/thread_cnt as usize)+1;
            for (chunk, chunk_src) in cubes[..].chunks_mut(chunk_len).zip(self.cubes[..].chunks(chunk_len)) {
                scope.spawn(move || {
                    for (c, cube) in chunk.iter_mut().zip(chunk_src) {
                        *c = CubeState {
                            pos: cube.pos + cube.spd * dt,
                            rot: Basis3::from_axis_angle(cube.rot_axe, Angle::from(deg(cube.rot_spd * dt))).concat(&cube.rot),
                            .. *cube
                        };
                    };
                });
            };
        });

//        for cube in &self.cubes {
//            let mut cs = CubeState {
//                pos: cube.pos + cube.spd * dt,
//                rot: Basis3::from_axis_angle(cube.rot_axe, Angle::from(deg(cube.rot_spd * dt))).concat(&cube.rot),
//                .. *cube};
//            cubes.push(cs);
//        }
        let (lx, ly) = f64::sin_cos(self.tick*0.1);
        let (lx, ly) = (lx as f32, ly as f32);
        let light_pos = v3(lx*50., ly*50., -3.);
        State {
            tick: self.tick + dt as f64,
            cubes: cubes,
            light_pos: light_pos,
        }
    }
}

use std::thread;
use std::sync::Arc;
use std::sync::mpsc::{sync_channel, channel, Receiver, SyncSender, Sender};

fn state_update_agent(thread_cnt: u32, rx: Receiver<(Arc<State>, f32)>, tx: Sender<State>) {
    loop {
        if let Ok((state_ref, dt)) = rx.recv() {
                if let Err(_) = tx.send(state_ref.update(dt, thread_cnt)) {
                    return ();
                }
        } else {
            return ();
        }
    }
}

pub struct StateUpdateAgent {
    tx: Sender<(Arc<State>, f32)>,
    rx: Receiver<State>,
}

impl StateUpdateAgent {
    pub fn new(thread_cnt: u32) -> Self {
        let (otx, orx) = channel();
        let (itx, irx) = channel();
        thread::spawn(move || state_update_agent(thread_cnt, orx, itx));
        StateUpdateAgent {
            tx: otx,
            rx: irx,
        }
    }

    pub fn start_update<'a>(&'a self, state: Arc<State>, dt: f32) -> UpdateResult<'a> {
        self.tx.send((state, dt)).unwrap();
        UpdateResult{
            agent: self,
            result_aquired: false,
        }
    }
}

pub struct UpdateResult<'a> {
    agent: &'a StateUpdateAgent,
    result_aquired: bool,
}

impl<'a> UpdateResult<'a> {
    pub fn get(&mut self) -> State {
        self.result_aquired = true;
        self.agent.rx.recv().unwrap()
    }
}

impl<'a> Drop for UpdateResult<'a> {
    fn drop(&mut self) {
        if !self.result_aquired {
            error!("Update result is not retrieved");
        }
    }
}
<commit_msg>Minor changes<commit_after>use cgmath::*;
use crossbeam;
use rand;
use rand::Rng;
use utils::v3;

pub type V3 = Vector3<f32>;

pub struct CubeState {
    pub pos: V3,
    pub rot: Basis3<f32>,
    pub spd: V3,
    pub rot_axe: V3,
    pub rot_spd: f32,
}

pub struct State {
    pub tick: f64,
    pub cubes: Vec<CubeState>,
    pub light_pos: V3,
}

const CUBE_SPAN : f32 = 200.;

impl State {
    pub fn new(object_count: u32) -> Self {
        let mut rng = rand::XorShiftRng::new_unseeded();
        let mut cubes = Vec::with_capacity(object_count as usize);
        for _ in 0 .. object_count {
            let cube_state = CubeState {
                pos: v3((rng.next_f32() - 0.5) * CUBE_SPAN,
                             (rng.next_f32() - 0.5) * CUBE_SPAN,
                             (rng.next_f32() - 1.) * CUBE_SPAN),
                rot: <Basis3<f32> as Rotation<_>>::one(),
                spd: v3(rng.next_f32() - 0.5,
                             rng.next_f32() - 0.5,
                             rng.next_f32() - 0.5).normalize()*rng.next_f32()*1.,
                rot_axe: v3(rng.next_f32() - 0.5,
                             rng.next_f32() - 0.5,
                             rng.next_f32() - 0.5)
                              .normalize(),
                rot_spd: rng.next_f32()*30.0,
            };
            cubes.push(cube_state);
        };
        State {
            tick: 0.0,
            cubes: cubes,
            light_pos: v3(0., 0., 0.),
        }
    }

    pub fn update(&self, dt: f32, thread_cnt: u32) -> Self {
        // advance position of cubes
        let mut cubes = Vec::with_capacity(self.cubes.len());
        // Uninitialized content of cubes will be immediately overwritten
        unsafe{ cubes.set_len(self.cubes.len()) };
        crossbeam::scope(|scope| {
            let chunk_len = (cubes.len()/thread_cnt as usize)+1;
            for (chunk, chunk_src) in cubes[..].chunks_mut(chunk_len).zip(self.cubes[..].chunks(chunk_len)) {
                scope.spawn(move || {
                    for (c, cube) in chunk.iter_mut().zip(chunk_src) {
                        let mut new_spd = cube.spd;
                        let mut new_pos = cube.pos + cube.spd * dt;
                        if new_pos.x < -CUBE_SPAN/2. {
                            new_pos.x = -CUBE_SPAN/2.;
                            new_spd.x = -new_spd.x;
                        }
                        if new_pos.x > CUBE_SPAN/2. {
                            new_pos.x = CUBE_SPAN/2.;
                            new_spd.x = -new_spd.x;
                        }
                        if new_pos.y < -CUBE_SPAN/2. {
                            new_pos.y = -CUBE_SPAN/2.;
                            new_spd.y = -new_spd.y;
                        }
                        if new_pos.y > CUBE_SPAN/2. {
                            new_pos.y = CUBE_SPAN/2.;
                            new_spd.y = -new_spd.y;
                        }
                        if new_pos.z < -CUBE_SPAN {
                            new_pos.z = -CUBE_SPAN;
                            new_spd.z = -new_spd.z;
                        }
                        if new_pos.z > 0. {
                            new_pos.z = 0.;
                            new_spd.z = -new_spd.z;
                        }
                        *c = CubeState {
                            pos: new_pos,
                            spd: new_spd,
                            rot: Basis3::from_axis_angle(cube.rot_axe, Angle::from(deg(cube.rot_spd * dt))).concat(&cube.rot),
                            .. *cube
                        };
                    };
                });
            };
        });

//        for cube in &self.cubes {
//            let mut cs = CubeState {
//                pos: cube.pos + cube.spd * dt,
//                rot: Basis3::from_axis_angle(cube.rot_axe, Angle::from(deg(cube.rot_spd * dt))).concat(&cube.rot),
//                .. *cube};
//            cubes.push(cs);
//        }
        let (lx, ly) = f64::sin_cos(self.tick*0.1);
        let (lx, ly) = (lx as f32, ly as f32);
        let light_pos = v3(lx*50., ly*50., -3.);
        State {
            tick: self.tick + dt as f64,
            cubes: cubes,
            light_pos: light_pos,
        }
    }
}

use std::thread;
use std::sync::Arc;
use std::sync::mpsc::{sync_channel, channel, Receiver, SyncSender, Sender};

fn state_update_agent(thread_cnt: u32, rx: Receiver<(Arc<State>, f32)>, tx: Sender<State>) {
    loop {
        if let Ok((state_ref, dt)) = rx.recv() {
            if let Err(_) = tx.send(state_ref.update(dt, thread_cnt)) {
                return ();
            }
        } else {
            return ();
        }
    }
}

pub struct StateUpdateAgent {
    tx: Sender<(Arc<State>, f32)>,
    rx: Receiver<State>,
}

impl StateUpdateAgent {
    pub fn new(thread_cnt: u32) -> Self {
        let (otx, orx) = channel();
        let (itx, irx) = channel();
        thread::spawn(move || state_update_agent(thread_cnt, orx, itx));
        StateUpdateAgent {
            tx: otx,
            rx: irx,
        }
    }

    pub fn start_update<'a>(&'a self, state: Arc<State>, dt: f32) -> UpdateResult<'a> {
        self.tx.send((state, dt)).unwrap();
        UpdateResult{
            agent: self,
            result_aquired: false,
        }
    }
}

pub struct UpdateResult<'a> {
    agent: &'a StateUpdateAgent,
    result_aquired: bool,
}

impl<'a> UpdateResult<'a> {
    pub fn get(&mut self) -> State {
        self.result_aquired = true;
        self.agent.rx.recv().unwrap()
    }
}

impl<'a> Drop for UpdateResult<'a> {
    fn drop(&mut self) {
        if !self.result_aquired {
            error!("Update result is not retrieved before drop");
        }
    }
}
<|endoftext|>


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

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

use ipnet::{IpNet, Ipv4Net, Ipv6Net};
use iprange::IpRange;
use regex::{RegexSet, RegexSetBuilder};

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

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

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

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

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

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

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

        f.write_str("] }")
    }
}

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

        Rules { ipv4, ipv6, rule }
    }

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

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

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

    /// Check if there are no rules for IPv4 addresses
    fn is_ipv4_empty(&self) -> bool {
        self.ipv4.iter().next().is_none()
    }

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

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

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

        let mut mode = Mode::BlackList;

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

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

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

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

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

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

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

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

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

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

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    pub fn check_host_in_proxy_list(&self, host: &str) -> Option<bool> {
        // Addresses in proxy_list will be proxied
        if self.white_list.check_host_matched(host) {
            return Some(true);
        }
        // Addresses in bypass_list will be bypassed
        if self.black_list.check_host_matched(host) {
            return Some(false);
        }
        None
    }

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

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

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

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

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

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

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

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

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

use ipnet::{IpNet, Ipv4Net, Ipv6Net};
use iprange::IpRange;
use regex::{RegexSet, RegexSetBuilder};

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

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

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

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

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

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

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

        f.write_str("] }")
    }
}

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

        Rules { ipv4, ipv6, rule }
    }

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

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

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

    /// Check if there are no rules for IPv4 addresses
    fn is_ipv4_empty(&self) -> bool {
        self.ipv4.iter().next().is_none()
    }

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

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

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

        let mut mode = Mode::BlackList;

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

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

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

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

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

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

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

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

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

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

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    pub fn check_host_in_proxy_list(&self, host: &str) -> Option<bool> {
        // Addresses in proxy_list will be proxied
        if self.white_list.check_host_matched(host) {
            return Some(true);
        }
        // Addresses in bypass_list will be bypassed
        if self.black_list.check_host_matched(host) {
            return Some(false);
        }
        None
    }

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

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

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

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

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

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

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

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


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use error::Result;
use error::GPSErrorKind as GPSEK;
use error::ResultExt;
use types::*;

use libimagstore::store::Entry;

use toml_query::read::TomlValueReadExt;
use toml_query::insert::TomlValueInsertExt;
use toml_query::delete::TomlValueDeleteExt;

pub trait GPSEntry {

    fn set_coordinates(&mut self, c: Coordinates) -> Result<()>;
    fn get_coordinates(&self) -> Result<Option<Coordinates>>;

    /// Remove the coordinates from the entry
    ///
    /// # Returns
    ///
    /// The return type is a bit complicated, but that has a reason:
    ///
    /// The outer Result<_> is used for notifying a failure during the header read/write action.
    /// If the Option<_> is Some(_), the value was deleted.
    /// The inner Result<_> is used for parsing failures during the parsing of the deleted value.
    ///
    /// So:
    ///
    /// * Ok(Some(Ok(_))) if the coordinates were deleted, returning the deleted value
    /// * Ok(Some(Err(_))) if the coordinates were deleted, but the deleted value couldn't be parsed
    /// * Ok(None) if there were no coordinates to delete
    /// * Err(e) if the deleting failed
    ///
    fn remove_coordinates(&mut self) -> Result<Option<Result<Coordinates>>>;

}

impl GPSEntry for Entry {

    fn set_coordinates(&mut self, c: Coordinates) -> Result<()> {
        self.get_header_mut()
            .insert("gps.coordinates", c.into())
            .map(|_| ())
            .chain_err(|| GPSEK::HeaderWriteError)
    }

    fn get_coordinates(&self) -> Result<Option<Coordinates>> {
        match self.get_header().read("gps.coordinates").chain_err(|| GPSEK::HeaderWriteError) {
            Ok(Some(hdr)) => Coordinates::from_value(hdr).map(Some),
            Ok(None)      => Ok(None),
            Err(e)        => Err(e),
        }
    }

    fn remove_coordinates(&mut self) -> Result<Option<Result<Coordinates>>> {
        let coordinates = self.get_coordinates();

        let patterns = [
            "gps.coordinates.latitude.degree",
            "gps.coordinates.latitude.minutes",
            "gps.coordinates.latitude.seconds",
            "gps.coordinates.longitude.degree",
            "gps.coordinates.longitude.minutes",
            "gps.coordinates.longitude.seconds",
            "gps.coordinates.latitude",
            "gps.coordinates.longitude",
            "gps.coordinates",
            "gps",
        ];

        let mut hdr = self.get_header_mut();
        for pattern in patterns.iter() {
            let _ = try!(hdr.delete(pattern).chain_err(|| GPSEK::HeaderWriteError));
        }

        match coordinates {
            Ok(None)       => Ok(None),
            Ok(Some(some)) => Ok(Some(Ok(some))),
            Err(e)         => Ok(Some(Err(e))),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use libimagstore::store::Store;

    use entry::*;

    fn setup_logging() {
        use env_logger;
        let _ = env_logger::init().unwrap_or(());
    }

    fn get_store() -> Store {
        use libimagstore::file_abstraction::InMemoryFileAbstraction;
        let backend = Box::new(InMemoryFileAbstraction::new());
        Store::new_with_backend(PathBuf::from("/"), None, backend).unwrap()
    }

    #[test]
    fn test_set_gps() {
        setup_logging();

        let store = get_store();

        let mut entry = store.create(PathBuf::from("test_set_gps")).unwrap();

        let coordinates = Coordinates {
            latitude: GPSValue::new(0, 0, 0),
            longitude: GPSValue::new(0, 0, 0),
        };

        let res = entry.set_coordinates(coordinates);

        assert!(res.is_ok());
    }

    #[test]
    fn test_setget_gps() {
        setup_logging();

        let store = get_store();

        let mut entry = store.create(PathBuf::from("test_setget_gps")).unwrap();

        let coordinates = Coordinates {
            latitude: GPSValue::new(0, 0, 0),
            longitude: GPSValue::new(0, 0, 0),
        };

        let res = entry.set_coordinates(coordinates);
        assert!(res.is_ok());

        let coordinates = entry.get_coordinates();

        assert!(coordinates.is_ok());
        let coordinates = coordinates.unwrap();

        assert!(coordinates.is_some());
        let coordinates = coordinates.unwrap();

        assert_eq!(0, coordinates.longitude.degree);
        assert_eq!(0, coordinates.longitude.minutes);
        assert_eq!(0, coordinates.longitude.seconds);
        assert_eq!(0, coordinates.latitude.degree);
        assert_eq!(0, coordinates.latitude.minutes);
        assert_eq!(0, coordinates.latitude.seconds);
    }
}

<commit_msg>Replace uses of try!() macro with "?" operator<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use error::Result;
use error::GPSErrorKind as GPSEK;
use error::ResultExt;
use types::*;

use libimagstore::store::Entry;

use toml_query::read::TomlValueReadExt;
use toml_query::insert::TomlValueInsertExt;
use toml_query::delete::TomlValueDeleteExt;

pub trait GPSEntry {

    fn set_coordinates(&mut self, c: Coordinates) -> Result<()>;
    fn get_coordinates(&self) -> Result<Option<Coordinates>>;

    /// Remove the coordinates from the entry
    ///
    /// # Returns
    ///
    /// The return type is a bit complicated, but that has a reason:
    ///
    /// The outer Result<_> is used for notifying a failure during the header read/write action.
    /// If the Option<_> is Some(_), the value was deleted.
    /// The inner Result<_> is used for parsing failures during the parsing of the deleted value.
    ///
    /// So:
    ///
    /// * Ok(Some(Ok(_))) if the coordinates were deleted, returning the deleted value
    /// * Ok(Some(Err(_))) if the coordinates were deleted, but the deleted value couldn't be parsed
    /// * Ok(None) if there were no coordinates to delete
    /// * Err(e) if the deleting failed
    ///
    fn remove_coordinates(&mut self) -> Result<Option<Result<Coordinates>>>;

}

impl GPSEntry for Entry {

    fn set_coordinates(&mut self, c: Coordinates) -> Result<()> {
        self.get_header_mut()
            .insert("gps.coordinates", c.into())
            .map(|_| ())
            .chain_err(|| GPSEK::HeaderWriteError)
    }

    fn get_coordinates(&self) -> Result<Option<Coordinates>> {
        match self.get_header().read("gps.coordinates").chain_err(|| GPSEK::HeaderWriteError) {
            Ok(Some(hdr)) => Coordinates::from_value(hdr).map(Some),
            Ok(None)      => Ok(None),
            Err(e)        => Err(e),
        }
    }

    fn remove_coordinates(&mut self) -> Result<Option<Result<Coordinates>>> {
        let coordinates = self.get_coordinates();

        let patterns = [
            "gps.coordinates.latitude.degree",
            "gps.coordinates.latitude.minutes",
            "gps.coordinates.latitude.seconds",
            "gps.coordinates.longitude.degree",
            "gps.coordinates.longitude.minutes",
            "gps.coordinates.longitude.seconds",
            "gps.coordinates.latitude",
            "gps.coordinates.longitude",
            "gps.coordinates",
            "gps",
        ];

        let mut hdr = self.get_header_mut();
        for pattern in patterns.iter() {
            let _ = hdr.delete(pattern).chain_err(|| GPSEK::HeaderWriteError)?;
        }

        match coordinates {
            Ok(None)       => Ok(None),
            Ok(Some(some)) => Ok(Some(Ok(some))),
            Err(e)         => Ok(Some(Err(e))),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use libimagstore::store::Store;

    use entry::*;

    fn setup_logging() {
        use env_logger;
        let _ = env_logger::init().unwrap_or(());
    }

    fn get_store() -> Store {
        use libimagstore::file_abstraction::InMemoryFileAbstraction;
        let backend = Box::new(InMemoryFileAbstraction::new());
        Store::new_with_backend(PathBuf::from("/"), None, backend).unwrap()
    }

    #[test]
    fn test_set_gps() {
        setup_logging();

        let store = get_store();

        let mut entry = store.create(PathBuf::from("test_set_gps")).unwrap();

        let coordinates = Coordinates {
            latitude: GPSValue::new(0, 0, 0),
            longitude: GPSValue::new(0, 0, 0),
        };

        let res = entry.set_coordinates(coordinates);

        assert!(res.is_ok());
    }

    #[test]
    fn test_setget_gps() {
        setup_logging();

        let store = get_store();

        let mut entry = store.create(PathBuf::from("test_setget_gps")).unwrap();

        let coordinates = Coordinates {
            latitude: GPSValue::new(0, 0, 0),
            longitude: GPSValue::new(0, 0, 0),
        };

        let res = entry.set_coordinates(coordinates);
        assert!(res.is_ok());

        let coordinates = entry.get_coordinates();

        assert!(coordinates.is_ok());
        let coordinates = coordinates.unwrap();

        assert!(coordinates.is_some());
        let coordinates = coordinates.unwrap();

        assert_eq!(0, coordinates.longitude.degree);
        assert_eq!(0, coordinates.longitude.minutes);
        assert_eq!(0, coordinates.longitude.seconds);
        assert_eq!(0, coordinates.latitude.degree);
        assert_eq!(0, coordinates.latitude.minutes);
        assert_eq!(0, coordinates.latitude.seconds);
    }
}

<|endoftext|>


<commit_before>use std::env;
use std::io::{stdout, stderr, Write};
use std::path::{Path, PathBuf};
use std::fs::metadata;
use std::os::unix::fs::MetadataExt;

use libc::getuid;

use options::build_mode::{build_mode, BuildMode};
use config::{Config, Settings, find_config};
use config::read_settings::{read_settings, MergedSettings};
use argparse::{ArgumentParser, Store, List, Collect, Print, StoreFalse, StoreTrue};

mod list;
mod user;
mod pack;
mod push;
pub mod build;
mod storage;
mod underscore;
mod completion;
mod commands;

mod supervisor;
mod simple;
mod capsule;

pub mod wrap;
pub mod system;
pub mod network;
mod socket;
mod volumes;
mod options;
mod prerequisites;

pub struct Context {
    pub config: Config,
    pub settings: Settings,
    pub ext_settings: MergedSettings,
    pub workdir: PathBuf,
    pub config_dir: PathBuf,
    pub build_mode: BuildMode,
    pub prerequisites: bool,
    pub isolate_network: bool
}

fn check_export(cmd: &String) -> Option<String> {
    if cmd == "/proc/self/exe" {
        return None;
    }
    match cmd.rfind("/") {
        Some(slash) => {
            let exe = &cmd[slash+1..];
            if exe == "vagga" || exe.starts_with("vagga_") {
                None
            } else {
                Some(exe.to_owned())
            }
        }
        None => Some(cmd.clone()),
    }
}

pub fn run(input_args: Vec<String>) -> i32 {
    let mut err = stderr();
    let workdir = env::current_dir().unwrap();

    let (config, cfg_dir) = match find_config(&workdir, true) {
        Ok(tup) => tup,
        Err(e) => {
            writeln!(&mut err, "{}", e).ok();
            return 126;
        }
    };
    // Not sure this is a best place, but the variable is needed for correct
    // reading of settings
    if let Some(x) = env::var_os("HOME") {
        env::set_var("_VAGGA_HOME", x);
    }

    let (ext_settings, int_settings) = match read_settings(&cfg_dir)
    {
        Ok(tup) => tup,
        Err(e) => {
            writeln!(&mut err, "{}", e).ok();
            return 126;
        }
    };

    let mut commands = Vec::<String>::new();
    let mut cname = "".to_string();
    let mut args = vec!();
    let mut set_env = Vec::<String>::new();
    let mut propagate_env = Vec::<String>::new();
    let mut bmode = Default::default();
    let mut owner_check = true;
    let mut prerequisites = true;
    let mut isolate_network = false;

    let export_command = input_args.get(0).and_then(check_export);
    if !int_settings.run_symlinks_as_commands || export_command.is_none() {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Runs a command in container, optionally builds container if that
            does not exists or outdated.

            Run `vagga` without arguments to see the list of commands.
            ");
        ap.add_option(&["-V", "--version"],
            Print(env!("VAGGA_VERSION").to_string()),
            "Show vagga version and exit");
        ap.refer(&mut set_env)
          .add_option(&["-E", "--env", "--environ"], Collect,
                "Set environment variable for running command")
          .metavar("NAME=VALUE");
        ap.refer(&mut propagate_env)
          .add_option(&["-e", "--use-env"], Collect,
                "Propagate variable VAR into command environment")
          .metavar("VAR");
        ap.refer(&mut owner_check)
          .add_option(&["--ignore-owner-check"], StoreFalse,
                "Ignore checking owner of the project directory");
        ap.refer(&mut prerequisites)
            .add_option(&["--no-prerequisites"], StoreFalse,
            "Run only specified command(s), don't run prerequisites");
        ap.refer(&mut isolate_network)
            .add_option(
                &["--isolate-network", "--no-network", "--no-net"],
                StoreTrue,
                "Run command(s) inside isolated network");
        build_mode(&mut ap, &mut bmode);
        ap.refer(&mut commands)
          .add_option(&["-m", "--run-multi"], List, "
            Run the following list of commands. Each without an arguments.
            When any of them fails, stop the chain. Basically it's the
            shortcut to `vagga cmd1 && vagga cmd2` except containers for
            `cmd2` are built beforehand, for your convenience. Also builtin
            commands (those starting with underscore) do not work with
            `vagga -m`");
        ap.refer(&mut cname)
          .add_argument("command", Store,
                "A vagga command to run");
        ap.refer(&mut args)
          .add_argument("args", List,
                "Arguments for the command");
        ap.stop_on_first_argument(true);
        ap.silence_double_dash(false);
        match ap.parse(input_args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return 0,
            Err(_) => return 122,
        }
    } else {
        let export_command = export_command.unwrap();
        for (name, cmd) in &config.commands {
            match cmd.link() {
                Some(ref lnk) if lnk.name == export_command => {
                    cname = name.clone();
                }
                _ => {}
            }
        }
        if cname == "" {
            writeln!(&mut err, "Can't find command {:?}", export_command).ok();
            return 127;
        }
        args = input_args;
        args.remove(0);  // the command itself
    }

    if &cname[..] == "_network" {
        args.insert(0, "vagga _network".to_string());
        return ::network::run(args);
    }

    if owner_check {
        let uid = unsafe { getuid() };
        match metadata(&cfg_dir) {
            Ok(ref stat) if stat.uid() == uid  => {}
            Ok(_) => {
                if uid == 0 {
                    writeln!(&mut err, "You should not run vagga as root \
                        (see http://bit.ly/err_root)").ok();
                    return 122;
                } else {
                    warn!("You are running vagga as a user \
                        different from the owner of project directory. \
                        You may not have needed permissions \
                        (see http://bit.ly/err_root)");
                }
            }
            Err(e) => {
                writeln!(&mut err, "Can't stat {:?}: {}", cfg_dir, e).ok();
                return 126;
            }
        }
    }

    let int_workdir = workdir.strip_prefix(&cfg_dir)
                             .unwrap_or(&Path::new("."));

    for k in propagate_env.into_iter() {
        if k.chars().find(|&c| c == '=').is_some() {
            writeln!(&mut err, "Environment variable name \
                (for option `-e`/`--use-env`) \
                can't contain equals `=` character. \
                To set key-value pair use `-E`/`--environ` option").ok();
            return 126;
        } else {
            env::set_var(&("VAGGAENV_".to_string() + &k[..]),
                env::var_os(&k).unwrap_or(From::from("")));
        }
    }
    for pair in set_env.into_iter() {
        let mut pairiter = pair[..].splitn(2, '=');
        let key = "VAGGAENV_".to_string() + pairiter.next().unwrap();
        if let Some(value) = pairiter.next() {
            env::set_var(&key, value.to_string());
        } else {
            env::remove_var(&key);
        }
    }

    let context = Context {
        config: config,
        settings: int_settings,
        ext_settings: ext_settings,
        workdir: int_workdir.to_path_buf(),
        config_dir: cfg_dir.to_path_buf(),
        build_mode: bmode,
        prerequisites: prerequisites,
        isolate_network: isolate_network,
    };

    if commands.len() > 0 {
        let result = user::run_multiple_commands(&context, commands);
        match result {
            Ok(rc) => {
                return rc;
            }
            Err(text) =>  {
                writeln!(&mut err, "{}", text).ok();
                return 121;
            }
        }
    }

    let result:Result<i32, String> = match &cname[..] {
        "" => {
            list::print_help(&context.config)
        }
        "_create_netns" => {
            network::create_netns(&context.config, args)
        }
        "_destroy_netns" => {
            network::destroy_netns(&context.config, args)
        }
        "_list" => {
            list::print_list(&context.config, args)
        }
        "_version_hash" => {
            underscore::version_hash(&context, &cname, args)
        }
        "_build_shell" | "_clean" | "_check_overlayfs_support" => {
            underscore::passthrough(&context, &cname, args)
        }
        "_base_dir" => {
            println!("{}", cfg_dir.display());
            Ok(0)
        }
        "_relative_work_dir" => {
            println!("{}", int_workdir.display());
            Ok(0)
        }
        "_build" => {
            build::build_command(&context, args)
        }
        "_run" => {
            underscore::run_command(&context, args)
        }
        "_run_in_netns" => {
            underscore::run_in_netns(&context, cname, args)
        }
        "_pack_image" => {
            pack::pack_command(&context, args)
        }
        "_push_image" => {
            push::push_command(&context, args)
        }
        "_init_storage_dir" => {
            storage::init_dir(&context.ext_settings, &cfg_dir, args)
        }
        "_compgen" => {
            completion::generate_completions(&context.config, args)
        }
        "_help" => {
            completion::generate_command_help(&context, args)
        }
        "_update_symlinks" => {
            commands::update_symlinks(&context, args)
        }
        "_capsule" => {
            ::capsule::run_command(&context, args)
        }
        _ => {
            user::run_user_command(&context, cname, args)
        }
    };

    match result {
        Ok(rc) => {
            return rc;
        }
        Err(text) =>  {
            writeln!(&mut err, "{}", text).ok();
            return 121;
        }
    }
}
<commit_msg>Fixes calling vagga as bare vagga<commit_after>use std::env;
use std::io::{stdout, stderr, Write};
use std::path::{Path, PathBuf};
use std::fs::metadata;
use std::os::unix::fs::MetadataExt;

use libc::getuid;

use options::build_mode::{build_mode, BuildMode};
use config::{Config, Settings, find_config};
use config::read_settings::{read_settings, MergedSettings};
use argparse::{ArgumentParser, Store, List, Collect, Print, StoreFalse, StoreTrue};

mod list;
mod user;
mod pack;
mod push;
pub mod build;
mod storage;
mod underscore;
mod completion;
mod commands;

mod supervisor;
mod simple;
mod capsule;

pub mod wrap;
pub mod system;
pub mod network;
mod socket;
mod volumes;
mod options;
mod prerequisites;

pub struct Context {
    pub config: Config,
    pub settings: Settings,
    pub ext_settings: MergedSettings,
    pub workdir: PathBuf,
    pub config_dir: PathBuf,
    pub build_mode: BuildMode,
    pub prerequisites: bool,
    pub isolate_network: bool
}

fn check_export(cmd: &String) -> Option<String> {
    if cmd == "/proc/self/exe" || cmd == "vagga" {
        return None;
    }
    match cmd.rfind("/") {
        Some(slash) => {
            let exe = &cmd[slash+1..];
            if exe == "vagga" || exe.starts_with("vagga_") {
                None
            } else {
                Some(exe.to_owned())
            }
        }
        None => Some(cmd.clone()),
    }
}

pub fn run(input_args: Vec<String>) -> i32 {
    let mut err = stderr();
    let workdir = env::current_dir().unwrap();

    let (config, cfg_dir) = match find_config(&workdir, true) {
        Ok(tup) => tup,
        Err(e) => {
            writeln!(&mut err, "{}", e).ok();
            return 126;
        }
    };
    // Not sure this is a best place, but the variable is needed for correct
    // reading of settings
    if let Some(x) = env::var_os("HOME") {
        env::set_var("_VAGGA_HOME", x);
    }

    let (ext_settings, int_settings) = match read_settings(&cfg_dir)
    {
        Ok(tup) => tup,
        Err(e) => {
            writeln!(&mut err, "{}", e).ok();
            return 126;
        }
    };

    let mut commands = Vec::<String>::new();
    let mut cname = "".to_string();
    let mut args = vec!();
    let mut set_env = Vec::<String>::new();
    let mut propagate_env = Vec::<String>::new();
    let mut bmode = Default::default();
    let mut owner_check = true;
    let mut prerequisites = true;
    let mut isolate_network = false;

    let export_command = input_args.get(0).and_then(check_export);
    if !int_settings.run_symlinks_as_commands || export_command.is_none() {
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Runs a command in container, optionally builds container if that
            does not exists or outdated.

            Run `vagga` without arguments to see the list of commands.
            ");
        ap.add_option(&["-V", "--version"],
            Print(env!("VAGGA_VERSION").to_string()),
            "Show vagga version and exit");
        ap.refer(&mut set_env)
          .add_option(&["-E", "--env", "--environ"], Collect,
                "Set environment variable for running command")
          .metavar("NAME=VALUE");
        ap.refer(&mut propagate_env)
          .add_option(&["-e", "--use-env"], Collect,
                "Propagate variable VAR into command environment")
          .metavar("VAR");
        ap.refer(&mut owner_check)
          .add_option(&["--ignore-owner-check"], StoreFalse,
                "Ignore checking owner of the project directory");
        ap.refer(&mut prerequisites)
            .add_option(&["--no-prerequisites"], StoreFalse,
            "Run only specified command(s), don't run prerequisites");
        ap.refer(&mut isolate_network)
            .add_option(
                &["--isolate-network", "--no-network", "--no-net"],
                StoreTrue,
                "Run command(s) inside isolated network");
        build_mode(&mut ap, &mut bmode);
        ap.refer(&mut commands)
          .add_option(&["-m", "--run-multi"], List, "
            Run the following list of commands. Each without an arguments.
            When any of them fails, stop the chain. Basically it's the
            shortcut to `vagga cmd1 && vagga cmd2` except containers for
            `cmd2` are built beforehand, for your convenience. Also builtin
            commands (those starting with underscore) do not work with
            `vagga -m`");
        ap.refer(&mut cname)
          .add_argument("command", Store,
                "A vagga command to run");
        ap.refer(&mut args)
          .add_argument("args", List,
                "Arguments for the command");
        ap.stop_on_first_argument(true);
        ap.silence_double_dash(false);
        match ap.parse(input_args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return 0,
            Err(_) => return 122,
        }
    } else {
        let export_command = export_command.unwrap();
        for (name, cmd) in &config.commands {
            match cmd.link() {
                Some(ref lnk) if lnk.name == export_command => {
                    cname = name.clone();
                }
                _ => {}
            }
        }
        if cname == "" {
            writeln!(&mut err, "Can't find command {:?}", export_command).ok();
            return 127;
        }
        args = input_args;
        args.remove(0);  // the command itself
    }

    if &cname[..] == "_network" {
        args.insert(0, "vagga _network".to_string());
        return ::network::run(args);
    }

    if owner_check {
        let uid = unsafe { getuid() };
        match metadata(&cfg_dir) {
            Ok(ref stat) if stat.uid() == uid  => {}
            Ok(_) => {
                if uid == 0 {
                    writeln!(&mut err, "You should not run vagga as root \
                        (see http://bit.ly/err_root)").ok();
                    return 122;
                } else {
                    warn!("You are running vagga as a user \
                        different from the owner of project directory. \
                        You may not have needed permissions \
                        (see http://bit.ly/err_root)");
                }
            }
            Err(e) => {
                writeln!(&mut err, "Can't stat {:?}: {}", cfg_dir, e).ok();
                return 126;
            }
        }
    }

    let int_workdir = workdir.strip_prefix(&cfg_dir)
                             .unwrap_or(&Path::new("."));

    for k in propagate_env.into_iter() {
        if k.chars().find(|&c| c == '=').is_some() {
            writeln!(&mut err, "Environment variable name \
                (for option `-e`/`--use-env`) \
                can't contain equals `=` character. \
                To set key-value pair use `-E`/`--environ` option").ok();
            return 126;
        } else {
            env::set_var(&("VAGGAENV_".to_string() + &k[..]),
                env::var_os(&k).unwrap_or(From::from("")));
        }
    }
    for pair in set_env.into_iter() {
        let mut pairiter = pair[..].splitn(2, '=');
        let key = "VAGGAENV_".to_string() + pairiter.next().unwrap();
        if let Some(value) = pairiter.next() {
            env::set_var(&key, value.to_string());
        } else {
            env::remove_var(&key);
        }
    }

    let context = Context {
        config: config,
        settings: int_settings,
        ext_settings: ext_settings,
        workdir: int_workdir.to_path_buf(),
        config_dir: cfg_dir.to_path_buf(),
        build_mode: bmode,
        prerequisites: prerequisites,
        isolate_network: isolate_network,
    };

    if commands.len() > 0 {
        let result = user::run_multiple_commands(&context, commands);
        match result {
            Ok(rc) => {
                return rc;
            }
            Err(text) =>  {
                writeln!(&mut err, "{}", text).ok();
                return 121;
            }
        }
    }

    let result:Result<i32, String> = match &cname[..] {
        "" => {
            list::print_help(&context.config)
        }
        "_create_netns" => {
            network::create_netns(&context.config, args)
        }
        "_destroy_netns" => {
            network::destroy_netns(&context.config, args)
        }
        "_list" => {
            list::print_list(&context.config, args)
        }
        "_version_hash" => {
            underscore::version_hash(&context, &cname, args)
        }
        "_build_shell" | "_clean" | "_check_overlayfs_support" => {
            underscore::passthrough(&context, &cname, args)
        }
        "_base_dir" => {
            println!("{}", cfg_dir.display());
            Ok(0)
        }
        "_relative_work_dir" => {
            println!("{}", int_workdir.display());
            Ok(0)
        }
        "_build" => {
            build::build_command(&context, args)
        }
        "_run" => {
            underscore::run_command(&context, args)
        }
        "_run_in_netns" => {
            underscore::run_in_netns(&context, cname, args)
        }
        "_pack_image" => {
            pack::pack_command(&context, args)
        }
        "_push_image" => {
            push::push_command(&context, args)
        }
        "_init_storage_dir" => {
            storage::init_dir(&context.ext_settings, &cfg_dir, args)
        }
        "_compgen" => {
            completion::generate_completions(&context.config, args)
        }
        "_help" => {
            completion::generate_command_help(&context, args)
        }
        "_update_symlinks" => {
            commands::update_symlinks(&context, args)
        }
        "_capsule" => {
            ::capsule::run_command(&context, args)
        }
        _ => {
            user::run_user_command(&context, cname, args)
        }
    };

    match result {
        Ok(rc) => {
            return rc;
        }
        Err(text) =>  {
            writeln!(&mut err, "{}", text).ok();
            return 121;
        }
    }
}
<|endoftext|>


<commit_before>//! Randomization of big integers

use rand::prelude::*;
use rand::distributions::uniform::{SampleUniform, UniformSampler};

use BigInt;
use BigUint;
use Sign::*;

use big_digit::BigDigit;
use bigint::{magnitude, into_magnitude};

use traits::Zero;
use integer::Integer;

pub trait RandBigInt {
    /// Generate a random `BigUint` of the given bit size.
    fn gen_biguint(&mut self, bit_size: usize) -> BigUint;

    /// Generate a random BigInt of the given bit size.
    fn gen_bigint(&mut self, bit_size: usize) -> BigInt;

    /// Generate a random `BigUint` less than the given bound. Fails
    /// when the bound is zero.
    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;

    /// Generate a random `BigUint` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;

    /// Generate a random `BigInt` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
}

impl<R: Rng + ?Sized> RandBigInt for R {
    fn gen_biguint(&mut self, bit_size: usize) -> BigUint {
        use super::big_digit::BITS;
        let (digits, rem) = bit_size.div_rem(&BITS);
        let mut data = Vec::with_capacity(digits + 1);
        for _ in 0..digits {
            data.push(self.gen());
        }
        if rem > 0 {
            let final_digit: BigDigit = self.gen();
            data.push(final_digit >> (BITS - rem));
        }
        BigUint::new(data)
    }

    fn gen_bigint(&mut self, bit_size: usize) -> BigInt {
        loop {
            // Generate a random BigUint...
            let biguint = self.gen_biguint(bit_size);
            // ...and then randomly assign it a Sign...
            let sign = if biguint.is_zero() {
                // ...except that if the BigUint is zero, we need to try
                // again with probability 0.5. This is because otherwise,
                // the probability of generating a zero BigInt would be
                // double that of any other number.
                if self.gen() {
                    continue;
                } else {
                    NoSign
                }
            } else if self.gen() {
                Plus
            } else {
                Minus
            };
            return BigInt::from_biguint(sign, biguint);
        }
    }

    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
        assert!(!bound.is_zero());
        let bits = bound.bits();
        loop {
            let n = self.gen_biguint(bits);
            if n < *bound {
                return n;
            }
        }
    }

    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
        assert!(*lbound < *ubound);
        if lbound.is_zero() {
            self.gen_biguint_below(ubound)
        } else {
            lbound + self.gen_biguint_below(&(ubound - lbound))
        }
    }

    fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt {
        assert!(*lbound < *ubound);
        if lbound.is_zero() {
            BigInt::from(self.gen_biguint_below(magnitude(&ubound)))
        } else if ubound.is_zero() {
            lbound + BigInt::from(self.gen_biguint_below(magnitude(&lbound)))
        } else {
            let delta = ubound - lbound;
            lbound + BigInt::from(self.gen_biguint_below(magnitude(&delta)))
        }
    }
}

/// The back-end implementing rand's `UniformSampler` for `BigUint`.
#[derive(Clone, Debug)]
pub struct UniformBigUint {
    base: BigUint,
    len: BigUint,
}

impl UniformSampler for UniformBigUint {
    type X = BigUint;

    fn new(low: Self::X, high: Self::X) -> Self {
        assert!(low < high);
        UniformBigUint {
            len: high - &low,
            base: low,
        }
    }

    fn new_inclusive(low: Self::X, high: Self::X) -> Self {
        assert!(low <= high);
        Self::new(low, high + 1u32)
    }

    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Self::X {
        &self.base + rng.gen_biguint_below(&self.len)
    }

    fn sample_single<R: Rng + ?Sized>(low: Self::X, high: Self::X, rng: &mut R) -> Self::X {
        rng.gen_biguint_range(&low, &high)
    }
}

impl SampleUniform for BigUint {
    type Sampler = UniformBigUint;
}

/// The back-end implementing rand's `UniformSampler` for `BigInt`.
#[derive(Clone, Debug)]
pub struct UniformBigInt {
    base: BigInt,
    len: BigUint,
}

impl UniformSampler for UniformBigInt {
    type X = BigInt;

    fn new(low: Self::X, high: Self::X) -> Self {
        assert!(low < high);
        UniformBigInt {
            len: into_magnitude(high - &low),
            base: low,
        }
    }

    fn new_inclusive(low: Self::X, high: Self::X) -> Self {
        assert!(low <= high);
        Self::new(low, high + 1u32)
    }

    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Self::X {
        &self.base + BigInt::from(rng.gen_biguint_below(&self.len))
    }

    fn sample_single<R: Rng + ?Sized>(low: Self::X, high: Self::X, rng: &mut R) -> Self::X {
        rng.gen_bigint_range(&low, &high)
    }
}

impl SampleUniform for BigInt {
    type Sampler = UniformBigInt;
}

/// A random distribution for `BigUint` and `BigInt` values of a particular bit size.
#[derive(Clone, Copy, Debug)]
pub struct RandomBits {
    bits: usize,
}

impl RandomBits {
    #[inline]
    pub fn new(bits: usize) -> RandomBits {
        RandomBits { bits }
    }
}

impl Distribution<BigUint> for RandomBits {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> BigUint {
        rng.gen_biguint(self.bits)
    }
}

impl Distribution<BigInt> for RandomBits {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> BigInt {
        rng.gen_bigint(self.bits)
    }
}
<commit_msg>Inline more random helpers<commit_after>//! Randomization of big integers

use rand::prelude::*;
use rand::distributions::uniform::{SampleUniform, UniformSampler};

use BigInt;
use BigUint;
use Sign::*;

use big_digit::BigDigit;
use bigint::{magnitude, into_magnitude};

use traits::Zero;
use integer::Integer;

pub trait RandBigInt {
    /// Generate a random `BigUint` of the given bit size.
    fn gen_biguint(&mut self, bit_size: usize) -> BigUint;

    /// Generate a random BigInt of the given bit size.
    fn gen_bigint(&mut self, bit_size: usize) -> BigInt;

    /// Generate a random `BigUint` less than the given bound. Fails
    /// when the bound is zero.
    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;

    /// Generate a random `BigUint` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;

    /// Generate a random `BigInt` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
}

impl<R: Rng + ?Sized> RandBigInt for R {
    fn gen_biguint(&mut self, bit_size: usize) -> BigUint {
        use super::big_digit::BITS;
        let (digits, rem) = bit_size.div_rem(&BITS);
        let mut data = Vec::with_capacity(digits + 1);
        for _ in 0..digits {
            data.push(self.gen());
        }
        if rem > 0 {
            let final_digit: BigDigit = self.gen();
            data.push(final_digit >> (BITS - rem));
        }
        BigUint::new(data)
    }

    fn gen_bigint(&mut self, bit_size: usize) -> BigInt {
        loop {
            // Generate a random BigUint...
            let biguint = self.gen_biguint(bit_size);
            // ...and then randomly assign it a Sign...
            let sign = if biguint.is_zero() {
                // ...except that if the BigUint is zero, we need to try
                // again with probability 0.5. This is because otherwise,
                // the probability of generating a zero BigInt would be
                // double that of any other number.
                if self.gen() {
                    continue;
                } else {
                    NoSign
                }
            } else if self.gen() {
                Plus
            } else {
                Minus
            };
            return BigInt::from_biguint(sign, biguint);
        }
    }

    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
        assert!(!bound.is_zero());
        let bits = bound.bits();
        loop {
            let n = self.gen_biguint(bits);
            if n < *bound {
                return n;
            }
        }
    }

    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
        assert!(*lbound < *ubound);
        if lbound.is_zero() {
            self.gen_biguint_below(ubound)
        } else {
            lbound + self.gen_biguint_below(&(ubound - lbound))
        }
    }

    fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt {
        assert!(*lbound < *ubound);
        if lbound.is_zero() {
            BigInt::from(self.gen_biguint_below(magnitude(&ubound)))
        } else if ubound.is_zero() {
            lbound + BigInt::from(self.gen_biguint_below(magnitude(&lbound)))
        } else {
            let delta = ubound - lbound;
            lbound + BigInt::from(self.gen_biguint_below(magnitude(&delta)))
        }
    }
}

/// The back-end implementing rand's `UniformSampler` for `BigUint`.
#[derive(Clone, Debug)]
pub struct UniformBigUint {
    base: BigUint,
    len: BigUint,
}

impl UniformSampler for UniformBigUint {
    type X = BigUint;

    #[inline]
    fn new(low: Self::X, high: Self::X) -> Self {
        assert!(low < high);
        UniformBigUint {
            len: high - &low,
            base: low,
        }
    }

    #[inline]
    fn new_inclusive(low: Self::X, high: Self::X) -> Self {
        assert!(low <= high);
        Self::new(low, high + 1u32)
    }

    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Self::X {
        &self.base + rng.gen_biguint_below(&self.len)
    }

    #[inline]
    fn sample_single<R: Rng + ?Sized>(low: Self::X, high: Self::X, rng: &mut R) -> Self::X {
        rng.gen_biguint_range(&low, &high)
    }
}

impl SampleUniform for BigUint {
    type Sampler = UniformBigUint;
}

/// The back-end implementing rand's `UniformSampler` for `BigInt`.
#[derive(Clone, Debug)]
pub struct UniformBigInt {
    base: BigInt,
    len: BigUint,
}

impl UniformSampler for UniformBigInt {
    type X = BigInt;

    #[inline]
    fn new(low: Self::X, high: Self::X) -> Self {
        assert!(low < high);
        UniformBigInt {
            len: into_magnitude(high - &low),
            base: low,
        }
    }

    #[inline]
    fn new_inclusive(low: Self::X, high: Self::X) -> Self {
        assert!(low <= high);
        Self::new(low, high + 1u32)
    }

    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Self::X {
        &self.base + BigInt::from(rng.gen_biguint_below(&self.len))
    }

    #[inline]
    fn sample_single<R: Rng + ?Sized>(low: Self::X, high: Self::X, rng: &mut R) -> Self::X {
        rng.gen_bigint_range(&low, &high)
    }
}

impl SampleUniform for BigInt {
    type Sampler = UniformBigInt;
}

/// A random distribution for `BigUint` and `BigInt` values of a particular bit size.
#[derive(Clone, Copy, Debug)]
pub struct RandomBits {
    bits: usize,
}

impl RandomBits {
    #[inline]
    pub fn new(bits: usize) -> RandomBits {
        RandomBits { bits }
    }
}

impl Distribution<BigUint> for RandomBits {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> BigUint {
        rng.gen_biguint(self.bits)
    }
}

impl Distribution<BigInt> for RandomBits {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> BigInt {
        rng.gen_bigint(self.bits)
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate extra;

use std::cell::Cell; // Provide mutable fields in immutable structs
use std::env;
use std::error::Error;
use std::fs;
use std::io::{self, Read, Stderr, StdoutLock, Write};
use std::process::exit;
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{cat} */ r#"NAME
    cat - concatenate files and print on the standard output

SYNOPSIS
    cat [-h | --help] [-A | --show-all] [-b | --number-nonblank] [-e] [-E | --show-ends]
        [-n | --number] [-s | --squeeze-blank] [-t] [-T] FILES...

DESCRIPTION
    Concatenates all files to the standard output.

    If no file is given, or if FILE is '-', read from standard input.

OPTIONS
    -A
    --show-all
        equivalent to -v -E -T

    -b
    --number-nonblank
        number nonempty output lines, overriding -n

    -e
        equivalent to -v -E

    -E
    --show-ends
        display $ at the end of each line

    -n
    --number
        number all output lines

    -s
    --squeeze-blank
        supress repeated empty output lines

    -t
        equivalent to -v -T

    -T
    --show_tabs
        display TAB characters as ^I

    -v
    --show-nonprinting
        use caret (^) and M- notation, except for LFD and TAB.

    -h
    --help
        display this help and exit

AUTHOR
    Written by Michael Murphy.
"#; /* @MANEND */

struct Program {
    exit_status:      Cell<i32>,
    number:           bool,
    number_nonblank:  bool,
    show_ends:        bool,
    show_tabs:        bool,
    show_nonprinting: bool,
    squeeze_blank:    bool,
    paths:            Vec<String>,
}

impl Program {
    /// Initialize the program's arguments and flags.
    fn initialize(stdout: &mut StdoutLock, stderr: &mut Stderr) -> Program {
        let mut cat = Program {
            exit_status:      Cell::new(0i32),
            number:           false,
            number_nonblank:  false,
            show_ends:        false,
            show_tabs:        false,
            show_nonprinting: false,
            squeeze_blank:    false,
            paths:            Vec::new()
        };
        for arg in env::args().skip(1) {
            if arg.starts_with('-') {
                match arg.as_str() {
                    "-h" | "--help" => {
                        stdout.write(MAN_PAGE.as_bytes()).try(stderr);
                        stdout.flush().try(stderr);
                        exit(0);
                    }
                    "-A" | "--show-all" => {
                        cat.show_nonprinting = true;
                        cat.show_ends = true;
                        cat.show_tabs = true;
                    },
                    "-b" | "--number-nonblank" => {
                        cat.number_nonblank = true;
                        cat.number = false;
                    },
                    "-e" => {
                        cat.show_nonprinting = true;
                        cat.show_ends = true;
                    },
                    "-E" | "--show-ends" => cat.show_ends = true,
                    "-n" | "--number" => {
                        cat.number = true;
                        cat.number_nonblank = false;
                    },
                    "-s" | "--squeeze-blank" => cat.squeeze_blank = true,
                    "-t" => {
                        cat.show_nonprinting = true;
                        cat.show_tabs = true;
                    },
                    "-T" => cat.show_tabs = true,
                    "-v" | "--show-nonprinting" => {
                        cat.show_nonprinting = true;
                    },
                    _ => {
                        stderr.write(b"invalid option -- '").try(stderr);
                        stderr.write(arg.as_bytes()).try(stderr);
                        stderr.write(b"'\nTry 'cat --help' for more information.\n").try(stderr);
                        stderr.flush().try(stderr);
                        exit(1);
                    }
                }
            } else {
                cat.paths.push(arg);
            }
        }
        cat
    }

    /// Execute the parameters given to the program.
    fn and_execute(&self, stdout: &mut StdoutLock, stderr: &mut Stderr) -> i32 {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();

        if self.paths.len() == 0 {
            io::copy(&mut stdin, stdout).try(stderr);
        } else {
            let mut line_count = 1;
            for path in &self.paths {
                if self.number || self.number_nonblank || self.show_ends || self.show_tabs || self.squeeze_blank || self.show_nonprinting {
                    let file = match fs::File::open(&path) {
                        Ok(file) => file,
                        Err(message) => {
                            stderr.write(&path.as_bytes()).try(stderr);
                            stderr.write(b": ").try(stderr);
                            stderr.write(message.description().as_bytes()).try(stderr);
                            stderr.write(b"\n").try(stderr);
                            stderr.flush().try(stderr);
                            self.exit_status.set(1i32);;
                            continue
                        }
                    };
                    let mut character_count = 0;
                    let mut last_line_was_blank = false;

                    for byte in file.bytes().map(|x| x.unwrap()) {
                        if (self.number && character_count == 0) || (character_count == 0 && self.number_nonblank && byte != b'\n') {
                            stdout.write(b"     ").try(stderr);
                            stdout.write(line_count.to_string().as_bytes()).try(stderr);
                            stdout.write(b"  ").try(stderr);
                            line_count += 1;
                        }
                        match byte {
                            0...8 | 11...31 => if self.show_nonprinting {
                                push_caret(stdout, stderr, byte+64);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                            9 => {
                                if self.show_tabs {
                                    push_caret(stdout, stderr, b'I');
                                } else {
                                    stdout.write(&[byte]).try(stderr);
                                }
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            }
                            10 => {
                                if character_count == 0 {
                                    if self.squeeze_blank && last_line_was_blank {
                                        continue
                                    } else if !last_line_was_blank {
                                        last_line_was_blank = true;
                                    }
                                } else {
                                    last_line_was_blank = false;
                                    character_count = 0;
                                }
                                if self.show_ends {
                                    stdout.write(b"$\n").try(stderr);
                                } else {
                                    stdout.write(b"\n").try(stderr);
                                }
                                stdout.flush().try(stderr);
                            },
                            32...126 => {
                                stdout.write(&[byte]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                            127 => if self.show_nonprinting {
                                push_caret(stdout, stderr, b'?');
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                            128...159 => if self.show_nonprinting {
                                stdout.write(b"M-^").try(stderr);
                                stdout.write(&[byte-64]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            } else {
                                stdout.write(&[byte]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                            _ => if self.show_nonprinting {
                                stdout.write(b"M-").try(stderr);
                                stdout.write(&[byte-128]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            } else {
                                stdout.write(&[byte]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                        }
                    }
                } else {
                    if path == "-" {
                        io::copy(&mut stdin, stdout).try(stderr);
                    } else {
                        let file = fs::File::open(&path);
                        let mut file = file.try(stderr);

                        io::copy(&mut file, stdout).try(stderr);
                    }
                }
            }
        }
        self.exit_status.get()
    }
}

/// Increase the character count by one if number printing is enabled.
fn count_character(character_count: &mut usize, number: &bool, number_nonblank: &bool) {
    if *number || *number_nonblank {
        *character_count += 1;
    }
}

/// Print a caret notation to stdout.
fn push_caret(stdout: &mut StdoutLock, stderr: &mut Stderr, notation: u8) {
    stdout.write(&[b'^']).try(stderr);
    stdout.write(&[notation]).try(stderr);
}

fn main() {
    let stdout = io::stdout();
    let mut stdout = stdout.lock();
    let mut stderr = io::stderr();
    exit(Program::initialize(&mut stdout, &mut stderr).and_execute(&mut stdout, &mut stderr));
}
<commit_msg>Increase Cat Performance With BufReader Iterate over a file using BufReader's bytes() implementation over Files. This leads to a huge performance increase. This makes it almost identical to the C version's performance.<commit_after>#![deny(warnings)]

extern crate extra;

use std::cell::Cell; // Provide mutable fields in immutable structs
use std::env;
use std::error::Error;
use std::fs;
use std::io::{self, BufReader, Read, Stderr, StdoutLock, Write};
use std::process::exit;
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{cat} */ r#"NAME
    cat - concatenate files and print on the standard output

SYNOPSIS
    cat [-h | --help] [-A | --show-all] [-b | --number-nonblank] [-e] [-E | --show-ends]
        [-n | --number] [-s | --squeeze-blank] [-t] [-T] FILES...

DESCRIPTION
    Concatenates all files to the standard output.

    If no file is given, or if FILE is '-', read from standard input.

OPTIONS
    -A
    --show-all
        equivalent to -v -E -T

    -b
    --number-nonblank
        number nonempty output lines, overriding -n

    -e
        equivalent to -v -E

    -E
    --show-ends
        display $ at the end of each line

    -n
    --number
        number all output lines

    -s
    --squeeze-blank
        supress repeated empty output lines

    -t
        equivalent to -v -T

    -T
    --show_tabs
        display TAB characters as ^I

    -v
    --show-nonprinting
        use caret (^) and M- notation, except for LFD and TAB.

    -h
    --help
        display this help and exit

AUTHOR
    Written by Michael Murphy.
"#; /* @MANEND */

struct Program {
    exit_status:      Cell<i32>,
    number:           bool,
    number_nonblank:  bool,
    show_ends:        bool,
    show_tabs:        bool,
    show_nonprinting: bool,
    squeeze_blank:    bool,
    paths:            Vec<String>,
}

impl Program {
    /// Initialize the program's arguments and flags.
    fn initialize(stdout: &mut StdoutLock, stderr: &mut Stderr) -> Program {
        let mut cat = Program {
            exit_status:      Cell::new(0i32),
            number:           false,
            number_nonblank:  false,
            show_ends:        false,
            show_tabs:        false,
            show_nonprinting: false,
            squeeze_blank:    false,
            paths:            Vec::new()
        };
        for arg in env::args().skip(1) {
            if arg.starts_with('-') {
                match arg.as_str() {
                    "-h" | "--help" => {
                        stdout.write(MAN_PAGE.as_bytes()).try(stderr);
                        stdout.flush().try(stderr);
                        exit(0);
                    }
                    "-A" | "--show-all" => {
                        cat.show_nonprinting = true;
                        cat.show_ends = true;
                        cat.show_tabs = true;
                    },
                    "-b" | "--number-nonblank" => {
                        cat.number_nonblank = true;
                        cat.number = false;
                    },
                    "-e" => {
                        cat.show_nonprinting = true;
                        cat.show_ends = true;
                    },
                    "-E" | "--show-ends" => cat.show_ends = true,
                    "-n" | "--number" => {
                        cat.number = true;
                        cat.number_nonblank = false;
                    },
                    "-s" | "--squeeze-blank" => cat.squeeze_blank = true,
                    "-t" => {
                        cat.show_nonprinting = true;
                        cat.show_tabs = true;
                    },
                    "-T" => cat.show_tabs = true,
                    "-v" | "--show-nonprinting" => {
                        cat.show_nonprinting = true;
                    },
                    _ => {
                        stderr.write(b"invalid option -- '").try(stderr);
                        stderr.write(arg.as_bytes()).try(stderr);
                        stderr.write(b"'\nTry 'cat --help' for more information.\n").try(stderr);
                        stderr.flush().try(stderr);
                        exit(1);
                    }
                }
            } else {
                cat.paths.push(arg);
            }
        }
        cat
    }

    /// Execute the parameters given to the program.
    fn and_execute(&self, stdout: &mut StdoutLock, stderr: &mut Stderr) -> i32 {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();

        if self.paths.len() == 0 {
            io::copy(&mut stdin, stdout).try(stderr);
        } else {
            let mut line_count = 1;
            for path in &self.paths {
                if self.number || self.number_nonblank || self.show_ends || self.show_tabs || self.squeeze_blank || self.show_nonprinting {
                    let file = match fs::File::open(&path) {
                        Ok(file) => file,
                        Err(message) => {
                            stderr.write(&path.as_bytes()).try(stderr);
                            stderr.write(b": ").try(stderr);
                            stderr.write(message.description().as_bytes()).try(stderr);
                            stderr.write(b"\n").try(stderr);
                            stderr.flush().try(stderr);
                            self.exit_status.set(1i32);;
                            continue
                        }
                    };
                    let mut character_count = 0;
                    let mut last_line_was_blank = false;

                    for byte in BufReader::new(&file).bytes().map(|x| x.unwrap()) {
                        if (self.number && character_count == 0) || (character_count == 0 && self.number_nonblank && byte != b'\n') {
                            stdout.write(b"     ").try(stderr);
                            stdout.write(line_count.to_string().as_bytes()).try(stderr);
                            stdout.write(b"  ").try(stderr);
                            line_count += 1;
                        }
                        match byte {
                            0...8 | 11...31 => if self.show_nonprinting {
                                push_caret(stdout, stderr, byte+64);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                            9 => {
                                if self.show_tabs {
                                    push_caret(stdout, stderr, b'I');
                                } else {
                                    stdout.write(&[byte]).try(stderr);
                                }
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            }
                            10 => {
                                if character_count == 0 {
                                    if self.squeeze_blank && last_line_was_blank {
                                        continue
                                    } else if !last_line_was_blank {
                                        last_line_was_blank = true;
                                    }
                                } else {
                                    last_line_was_blank = false;
                                    character_count = 0;
                                }
                                if self.show_ends {
                                    stdout.write(b"$\n").try(stderr);
                                } else {
                                    stdout.write(b"\n").try(stderr);
                                }
                                stdout.flush().try(stderr);
                            },
                            32...126 => {
                                stdout.write(&[byte]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                            127 => if self.show_nonprinting {
                                push_caret(stdout, stderr, b'?');
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                            128...159 => if self.show_nonprinting {
                                stdout.write(b"M-^").try(stderr);
                                stdout.write(&[byte-64]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            } else {
                                stdout.write(&[byte]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                            _ => if self.show_nonprinting {
                                stdout.write(b"M-").try(stderr);
                                stdout.write(&[byte-128]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            } else {
                                stdout.write(&[byte]).try(stderr);
                                count_character(&mut character_count, &self.number, &self.number_nonblank);
                            },
                        }
                    }
                } else {
                    if path == "-" {
                        io::copy(&mut stdin, stdout).try(stderr);
                    } else {
                        let file = fs::File::open(&path);
                        let mut file = file.try(stderr);

                        io::copy(&mut file, stdout).try(stderr);
                    }
                }
            }
        }
        self.exit_status.get()
    }
}

/// Increase the character count by one if number printing is enabled.
fn count_character(character_count: &mut usize, number: &bool, number_nonblank: &bool) {
    if *number || *number_nonblank {
        *character_count += 1;
    }
}

/// Print a caret notation to stdout.
fn push_caret(stdout: &mut StdoutLock, stderr: &mut Stderr, notation: u8) {
    stdout.write(&[b'^']).try(stderr);
    stdout.write(&[notation]).try(stderr);
}

fn main() {
    let stdout = io::stdout();
    let mut stdout = stdout.lock();
    let mut stderr = io::stderr();
    exit(Program::initialize(&mut stdout, &mut stderr).and_execute(&mut stdout, &mut stderr));
}
<|endoftext|>