<commit_before>use std;
use std::fmt;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash, PartialOrd, Ord)]
pub enum IpNet {
    V4(Ipv4Net),
    V6(Ipv6Net),
}

#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash, PartialOrd, Ord)]
pub struct Ipv4Net {
    addr: Ipv4Addr,
    prefix_len: u8,
}

#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash, PartialOrd, Ord)]
pub struct Ipv6Net {
    addr: Ipv6Addr,
    prefix_len: u8,
}

// For the time being deref method calls to the IpAddr implemenations.

/*impl std::ops::Deref for IpNet {
    type Target = IpAddr;
    fn deref(&self) -> &Self::Target {
        match *self {
            IpNet::V4(ref a) => a.addr,
            IpNet::V6(ref a) => a.addr,
        }
    }
}*/

impl std::ops::Deref for Ipv4Net {
    type Target = Ipv4Addr;
    fn deref(&self) -> &Self::Target {
        &self.addr
    }
}

impl std::ops::Deref for Ipv6Net {
    type Target = Ipv6Addr;
    fn deref(&self) -> &Self::Target {
        &self.addr
    }
}

impl IpNet {
    pub fn supernet(&self) -> IpNet {
        match *self {
            IpNet::V4(ref a) => IpNet::V4(a.supernet()),
            IpNet::V6(ref a) => IpNet::V6(a.supernet()),
        }
    }

    pub fn contains(&self, other: &IpNet) -> bool {
        match (*self, *other) {
            (IpNet::V4(ref a), IpNet::V4(ref b)) => a.contains(b),
            (IpNet::V6(ref a), IpNet::V6(ref b)) => a.contains(b),
            (_, _) => false,
        }
    }

    pub fn sibling(&self, other: &IpNet) -> bool {
        match (*self, *other) {
            (IpNet::V4(ref a), IpNet::V4(ref b)) => a.sibling(b),
            (IpNet::V6(ref a), IpNet::V6(ref b)) => a.sibling(b),
            (_, _) => false,
        }
    }
}

impl Ipv4Net {
    pub fn new(ip: Ipv4Addr, prefix_len: u8) -> Ipv4Net {
        // TODO: Should error if prefix_len > 32 as prefix_len <= 32 is
        // assumed in subsequent methods.
        Ipv4Net { addr: ip, prefix_len: prefix_len, }
    }

    pub fn netmask(&self) -> Ipv4Addr {
        Ipv4Addr::from(
            // Avoid deny(exceeding_bitshifts)
            if self.prefix_len > 0 {
                0xffffffff << 32 - self.prefix_len
            }
            else {
                0x00000000
            }
        )
    }

    pub fn hostmask(&self) -> Ipv4Addr {
        Ipv4Addr::from(
            // Avoid deny(exceeding_bitshifts)
            if self.prefix_len < 32 {
                0xffffffff >> self.prefix_len
            }
            else {
                0x00000000
            }
        )
    }

    pub fn network(&self) -> Ipv4Addr {
        // BitAnd is not implemented for Ipv4Addr.
        Ipv4Addr::from(u32::from(self.addr) & u32::from(self.netmask()))
    }

    pub fn broadcast(&self) -> Ipv4Addr {
        // BitOr is not implemented for Ipv4Addr.
        Ipv4Addr::from(u32::from(self.addr) | u32::from(self.hostmask()))
    }

    pub fn supernet(&self) -> Ipv4Net {
        Ipv4Net::new(self.addr.clone(), self.prefix_len - 1)
    }

    // TODO: Should the parameter be depth, as in the number of
    // divisions to make, or new_prefix_len (self explanatory), or
    // num_subnets? I suspect in use the later two will be more useful.
    // TODO: Should this be an iterator?
    pub fn subnets(&self, new_prefix_len: u8) -> Vec<Ipv4Net> {
        // TODO: Need to test bounds of new_prefix_len and error.
        // Or perhaps just set to 32 in that case.
        // TODO: std::ops::Add missing for Ipv4Addr
        let rng = 2u32.pow(32 - new_prefix_len as u32);
        let mut net = u32::from(self.network());
        let end = u32::from(self.broadcast());
        let mut res: Vec<Ipv4Net> = Vec::new();

        while net < end {
            res.push(Ipv4Net::new(Ipv4Addr::from(net), new_prefix_len));
            net += rng;
        }
        res

    }

    pub fn contains(&self, other: &Ipv4Net) -> bool {
        self.network() <= other.network() && self.broadcast() >= other.broadcast()
    }

    pub fn sibling(&self, other: &Ipv4Net) -> bool {
        self.prefix_len == other.prefix_len && self.supernet().contains(other)
    }
}

impl Ipv6Net {
    pub fn new(ip: Ipv6Addr, prefix_len: u8) -> Ipv6Net {
        // TODO: Should error if prefix_len > 128 as prefix_len <= 128
        // is assumed in subsequent methods.
        Ipv6Net { addr: ip, prefix_len: prefix_len, }
    }

    pub fn addr(&self) -> Ipv6Addr {
        self.addr
    }

    // The u128 type would be nice here, but it's not marked stable yet.
    pub fn netmask(&self) -> Ipv6Addr {
        // Avoid deny(exceeding_bitshifts)
        let a: u64 = if self.prefix_len > 0 { 0xffff_ffff_ffff_ffff << 64u8.saturating_sub(self.prefix_len) } else { 0x0 };
        let b: u64 = if self.prefix_len > 64 { 0xffff_ffff_ffff_ffff << 128 - self.prefix_len } else { 0x0 };

        Ipv6Addr::new(
            (a >> 48) as u16, (a >> 32) as u16, (a >> 16) as u16, a as u16,
            (b >> 48) as u16, (b >> 32) as u16, (b >> 16) as u16, b as u16
        )
    }

    // The u128 type would be nice here, but it's not marked stable yet.
    pub fn hostmask(&self) -> Ipv6Addr {
        // Avoid deny(exceeding_bitshifts)
        let a: u64 = if self.prefix_len < 64 { 0xffff_ffff_ffff_ffff >> self.prefix_len } else { 0x0 };
        let b: u64 = if self.prefix_len < 128 { 0xffff_ffff_ffff_ffff >> self.prefix_len.saturating_sub(64) } else { 0x0 };

        Ipv6Addr::new(
            (a >> 48) as u16, (a >> 32) as u16, (a >> 16) as u16, a as u16,
            (b >> 48) as u16, (b >> 32) as u16, (b >> 16) as u16, b as u16
        )
    }

    // The u128 type would be nice here, but it's not marked stable yet.
    pub fn network(&self) -> Ipv6Addr {
        // BitAnd is not implemented for Ipv6Addr.
        let ip = self.segments();
        let m = self.netmask().segments();

        Ipv6Addr::new(
            ip[0] & m[0], ip[1] & m[1], ip[2] & m[2], ip[3] & m[3],
            ip[4] & m[4], ip[5] & m[5], ip[6] & m[6], ip[7] & m[7]
        )
    }

    // TODO: Technically there is no such thing as a broadcast. Perhaps
    // we should change the network() and broadcast() methods to be
    // first() and last() or similar.
    // The u128 type would be nice here, but it's not marked stable yet.
    pub fn broadcast(&self) -> Ipv6Addr {
        // BitOr is not implemented for Ipv4Addr.
        let ip = self.segments();
        let m = self.hostmask().segments();

        Ipv6Addr::new(
            ip[0] | m[0], ip[1] | m[1], ip[2] | m[2], ip[3] | m[3],
            ip[4] | m[4], ip[5] | m[5], ip[6] | m[6], ip[7] | m[7]
        )
    }

    pub fn supernet(&self) -> Ipv6Net {
        Ipv6Net::new(self.addr.clone(), self.prefix_len - 1)
    }

    // TODO: Is there a nicer way to do this? Using u128 types would be
    // nice but they're not marked as stable yet. Should this return an
    // iterator instead of a vector?
    pub fn subnets(&self, new_prefix_len: u8) -> Vec<Ipv6Net> {
        // TODO: Need to implement a proper error handling scheme.
        if new_prefix_len <= self.prefix_len {
            return Vec::new();
        }

        let mut network = ipv6_addr_into_double_u64(self.network());
        let mut broadcast = ipv6_addr_into_double_u64(self.broadcast());

        let mask = ipv6_addr_into_double_u64(Ipv6Net::new(self.addr.clone(), new_prefix_len).hostmask());
        let mut res: Vec<Ipv6Net> = Vec::new();

        while network < broadcast {
            res.push(Ipv6Net::new(ipv6_addr_from_double_u64(network), new_prefix_len));

            network[0] |= mask[0];
            network[1] |= mask[1];

            let (lo, ov) = network[1].overflowing_add(1);
            network[1] = lo;

            if ov {
                network[0] += 1;
            }
        }

        res
    }

    pub fn contains(&self, other: &Ipv6Net) -> bool {
        self.network() <= other.network() && self.broadcast() >= other.broadcast()
    }

    pub fn sibling(&self, other: &Ipv6Net) -> bool {
        self.prefix_len == other.prefix_len && self.supernet().contains(other)
    }
}

// TODO: It would be nice to implement From on Ipv6Addr for this.
pub fn ipv6_addr_from_double_u64(segments: [u64; 2]) -> Ipv6Addr {
    let (hi, lo) = (segments[0], segments[1]);
    Ipv6Addr::new(
        (hi >> 48) as u16, (hi >> 32) as u16, (hi >> 16) as u16, hi as u16,
        (lo >> 48) as u16, (lo >> 32) as u16, (lo >> 16) as u16, lo as u16
    )
}

// TODO: It would be nice to implement From on Ipv6Addr for this.
pub fn ipv6_addr_into_double_u64(ip: Ipv6Addr) -> [u64; 2] {
    let ip = ip.octets();
    [
        ((ip[0] as u64) << 56) + ((ip[1] as u64) << 48) +
        ((ip[2] as u64) << 40) + ((ip[3] as u64) << 32) +
        ((ip[4] as u64) << 24) + ((ip[5] as u64) << 16) +
        ((ip[6] as u64) << 8) + (ip[7] as u64),
        ((ip[8] as u64) << 56) + ((ip[9] as u64) << 48) +
        ((ip[10] as u64) << 40) + ((ip[11] as u64) << 32) +
        ((ip[12] as u64) << 24) + ((ip[13] as u64) << 16) +
        ((ip[14] as u64) << 8) + (ip[15] as u64),
    ]
}

impl fmt::Display for IpNet {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            IpNet::V4(ref a) => a.fmt(fmt),
            IpNet::V6(ref a) => a.fmt(fmt),
        }
    }
}

impl fmt::Display for Ipv4Net {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}/{}", self.addr, self.prefix_len)
    }
}

impl fmt::Display for Ipv6Net {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}/{}", self.addr, self.prefix_len)
    }
}
<commit_msg>Building mask straight from new_prefix_len in Ipv6Net.subnets() instead of building a Ipv6Net object and calling hostmask on it.<commit_after>use std;
use std::fmt;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash, PartialOrd, Ord)]
pub enum IpNet {
    V4(Ipv4Net),
    V6(Ipv6Net),
}

#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash, PartialOrd, Ord)]
pub struct Ipv4Net {
    addr: Ipv4Addr,
    prefix_len: u8,
}

#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash, PartialOrd, Ord)]
pub struct Ipv6Net {
    addr: Ipv6Addr,
    prefix_len: u8,
}

// For the time being deref method calls to the IpAddr implemenations.

/*impl std::ops::Deref for IpNet {
    type Target = IpAddr;
    fn deref(&self) -> &Self::Target {
        match *self {
            IpNet::V4(ref a) => a.addr,
            IpNet::V6(ref a) => a.addr,
        }
    }
}*/

impl std::ops::Deref for Ipv4Net {
    type Target = Ipv4Addr;
    fn deref(&self) -> &Self::Target {
        &self.addr
    }
}

impl std::ops::Deref for Ipv6Net {
    type Target = Ipv6Addr;
    fn deref(&self) -> &Self::Target {
        &self.addr
    }
}

impl IpNet {
    pub fn supernet(&self) -> IpNet {
        match *self {
            IpNet::V4(ref a) => IpNet::V4(a.supernet()),
            IpNet::V6(ref a) => IpNet::V6(a.supernet()),
        }
    }

    pub fn contains(&self, other: &IpNet) -> bool {
        match (*self, *other) {
            (IpNet::V4(ref a), IpNet::V4(ref b)) => a.contains(b),
            (IpNet::V6(ref a), IpNet::V6(ref b)) => a.contains(b),
            (_, _) => false,
        }
    }

    pub fn sibling(&self, other: &IpNet) -> bool {
        match (*self, *other) {
            (IpNet::V4(ref a), IpNet::V4(ref b)) => a.sibling(b),
            (IpNet::V6(ref a), IpNet::V6(ref b)) => a.sibling(b),
            (_, _) => false,
        }
    }
}

impl Ipv4Net {
    pub fn new(ip: Ipv4Addr, prefix_len: u8) -> Ipv4Net {
        // TODO: Should error if prefix_len > 32 as prefix_len <= 32 is
        // assumed in subsequent methods.
        Ipv4Net { addr: ip, prefix_len: prefix_len, }
    }

    pub fn netmask(&self) -> Ipv4Addr {
        Ipv4Addr::from(
            // Avoid deny(exceeding_bitshifts)
            if self.prefix_len > 0 {
                0xffffffff << 32 - self.prefix_len
            }
            else {
                0x00000000
            }
        )
    }

    pub fn hostmask(&self) -> Ipv4Addr {
        Ipv4Addr::from(
            // Avoid deny(exceeding_bitshifts)
            if self.prefix_len < 32 {
                0xffffffff >> self.prefix_len
            }
            else {
                0x00000000
            }
        )
    }

    pub fn network(&self) -> Ipv4Addr {
        // BitAnd is not implemented for Ipv4Addr.
        Ipv4Addr::from(u32::from(self.addr) & u32::from(self.netmask()))
    }

    pub fn broadcast(&self) -> Ipv4Addr {
        // BitOr is not implemented for Ipv4Addr.
        Ipv4Addr::from(u32::from(self.addr) | u32::from(self.hostmask()))
    }

    pub fn supernet(&self) -> Ipv4Net {
        Ipv4Net::new(self.addr.clone(), self.prefix_len - 1)
    }

    // TODO: Should the parameter be depth, as in the number of
    // divisions to make, or new_prefix_len (self explanatory), or
    // num_subnets? I suspect in use the later two will be more useful.
    // TODO: Should this be an iterator?
    pub fn subnets(&self, new_prefix_len: u8) -> Vec<Ipv4Net> {
        // TODO: Need to test bounds of new_prefix_len and error.
        // Or perhaps just set to 32 in that case.
        // TODO: std::ops::Add missing for Ipv4Addr
        let rng = 2u32.pow(32 - new_prefix_len as u32);
        let mut net = u32::from(self.network());
        let end = u32::from(self.broadcast());
        let mut res: Vec<Ipv4Net> = Vec::new();

        while net < end {
            res.push(Ipv4Net::new(Ipv4Addr::from(net), new_prefix_len));
            net += rng;
        }
        res

    }

    pub fn contains(&self, other: &Ipv4Net) -> bool {
        self.network() <= other.network() && self.broadcast() >= other.broadcast()
    }

    pub fn sibling(&self, other: &Ipv4Net) -> bool {
        self.prefix_len == other.prefix_len && self.supernet().contains(other)
    }
}

impl Ipv6Net {
    pub fn new(ip: Ipv6Addr, prefix_len: u8) -> Ipv6Net {
        // TODO: Should error if prefix_len > 128 as prefix_len <= 128
        // is assumed in subsequent methods.
        Ipv6Net { addr: ip, prefix_len: prefix_len, }
    }

    pub fn addr(&self) -> Ipv6Addr {
        self.addr
    }

    // The u128 type would be nice here, but it's not marked stable yet.
    pub fn netmask(&self) -> Ipv6Addr {
        // Avoid deny(exceeding_bitshifts)
        let a: u64 = if self.prefix_len > 0 { 0xffff_ffff_ffff_ffff << 64u8.saturating_sub(self.prefix_len) } else { 0x0 };
        let b: u64 = if self.prefix_len > 64 { 0xffff_ffff_ffff_ffff << 128 - self.prefix_len } else { 0x0 };

        Ipv6Addr::new(
            (a >> 48) as u16, (a >> 32) as u16, (a >> 16) as u16, a as u16,
            (b >> 48) as u16, (b >> 32) as u16, (b >> 16) as u16, b as u16
        )
    }

    // The u128 type would be nice here, but it's not marked stable yet.
    pub fn hostmask(&self) -> Ipv6Addr {
        // Avoid deny(exceeding_bitshifts)
        let a: u64 = if self.prefix_len < 64 { 0xffff_ffff_ffff_ffff >> self.prefix_len } else { 0x0 };
        let b: u64 = if self.prefix_len < 128 { 0xffff_ffff_ffff_ffff >> self.prefix_len.saturating_sub(64) } else { 0x0 };

        Ipv6Addr::new(
            (a >> 48) as u16, (a >> 32) as u16, (a >> 16) as u16, a as u16,
            (b >> 48) as u16, (b >> 32) as u16, (b >> 16) as u16, b as u16
        )
    }

    // The u128 type would be nice here, but it's not marked stable yet.
    pub fn network(&self) -> Ipv6Addr {
        // BitAnd is not implemented for Ipv6Addr.
        let ip = self.segments();
        let m = self.netmask().segments();

        Ipv6Addr::new(
            ip[0] & m[0], ip[1] & m[1], ip[2] & m[2], ip[3] & m[3],
            ip[4] & m[4], ip[5] & m[5], ip[6] & m[6], ip[7] & m[7]
        )
    }

    // TODO: Technically there is no such thing as a broadcast. Perhaps
    // we should change the network() and broadcast() methods to be
    // first() and last() or similar.
    // The u128 type would be nice here, but it's not marked stable yet.
    pub fn broadcast(&self) -> Ipv6Addr {
        // BitOr is not implemented for Ipv4Addr.
        let ip = self.segments();
        let m = self.hostmask().segments();

        Ipv6Addr::new(
            ip[0] | m[0], ip[1] | m[1], ip[2] | m[2], ip[3] | m[3],
            ip[4] | m[4], ip[5] | m[5], ip[6] | m[6], ip[7] | m[7]
        )
    }

    pub fn supernet(&self) -> Ipv6Net {
        Ipv6Net::new(self.addr.clone(), self.prefix_len - 1)
    }

    // TODO: Is there a nicer way to do this? Using u128 types would be
    // nice but they're not marked as stable yet. Should this return an
    // iterator instead of a vector?
    pub fn subnets(&self, new_prefix_len: u8) -> Vec<Ipv6Net> {
        // TODO: Need to implement a proper error handling scheme.
        if new_prefix_len <= self.prefix_len {
            return Vec::new();
        }

        let mut network = ipv6_addr_into_double_u64(self.network());
        let mut broadcast = ipv6_addr_into_double_u64(self.broadcast());

        let mask = ipv6_addr_into_double_u64(Ipv6Net::new(self.addr.clone(), new_prefix_len).hostmask());

        let mask: [u64; 2] = [
            if new_prefix_len < 64 { 0xffff_ffff_ffff_ffff >> new_prefix_len } else { 0x0 },
            if new_prefix_len >= 64 && new_prefix_len < 128 { 0xffff_ffff_ffff_ffff >> new_prefix_len - 64 } else { 0x0 }
        ];

        let mut res: Vec<Ipv6Net> = Vec::new();

        while network < broadcast {
            res.push(Ipv6Net::new(ipv6_addr_from_double_u64(network), new_prefix_len));

            network[0] |= mask[0];
            network[1] |= mask[1];

            let (lo, ov) = network[1].overflowing_add(1);
            network[1] = lo;

            if ov {
                network[0] += 1;
            }
        }

        res
    }

    pub fn contains(&self, other: &Ipv6Net) -> bool {
        self.network() <= other.network() && self.broadcast() >= other.broadcast()
    }

    pub fn sibling(&self, other: &Ipv6Net) -> bool {
        self.prefix_len == other.prefix_len && self.supernet().contains(other)
    }
}

// TODO: It would be nice to implement From on Ipv6Addr for this.
pub fn ipv6_addr_from_double_u64(segments: [u64; 2]) -> Ipv6Addr {
    let (hi, lo) = (segments[0], segments[1]);
    Ipv6Addr::new(
        (hi >> 48) as u16, (hi >> 32) as u16, (hi >> 16) as u16, hi as u16,
        (lo >> 48) as u16, (lo >> 32) as u16, (lo >> 16) as u16, lo as u16
    )
}

// TODO: It would be nice to implement From on Ipv6Addr for this.
pub fn ipv6_addr_into_double_u64(ip: Ipv6Addr) -> [u64; 2] {
    let ip = ip.octets();
    [
        ((ip[0] as u64) << 56) + ((ip[1] as u64) << 48) +
        ((ip[2] as u64) << 40) + ((ip[3] as u64) << 32) +
        ((ip[4] as u64) << 24) + ((ip[5] as u64) << 16) +
        ((ip[6] as u64) << 8) + (ip[7] as u64),
        ((ip[8] as u64) << 56) + ((ip[9] as u64) << 48) +
        ((ip[10] as u64) << 40) + ((ip[11] as u64) << 32) +
        ((ip[12] as u64) << 24) + ((ip[13] as u64) << 16) +
        ((ip[14] as u64) << 8) + (ip[15] as u64),
    ]
}

impl fmt::Display for IpNet {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            IpNet::V4(ref a) => a.fmt(fmt),
            IpNet::V6(ref a) => a.fmt(fmt),
        }
    }
}

impl fmt::Display for Ipv4Net {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}/{}", self.addr, self.prefix_len)
    }
}

impl fmt::Display for Ipv6Net {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}/{}", self.addr, self.prefix_len)
    }
}
<|endoftext|>


<commit_before>#![crate_name = "uutils"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Michael Gehring <mg@ebfe.org>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

include!(concat!(env!("OUT_DIR"), "/uutils_crates.rs"));

use std::collections::hash_map::HashMap;
use std::path::Path;
use std::env;
use std::io::Write;

extern crate uucore;

static NAME: &'static str = "uutils";
static VERSION: &'static str = env!("CARGO_PKG_VERSION");

include!(concat!(env!("OUT_DIR"), "/uutils_map.rs"));

fn usage(cmap: &UtilityMap) {
    println!("{} {}", NAME, VERSION);
    println!("");
    println!("Usage:");
    println!("  {} [util [arguments...]]\n", NAME);
    println!("Currently defined functions:");
    let mut utils: Vec<&str> = cmap.keys().map(|&s| s).collect();
    utils.sort();
    for util in utils {
        println!("\t{}", util);
    }
}

fn main() {
    uucore::panic::install_sigpipe_hook();

    let umap = util_map();
    let mut args: Vec<String> = env::args().collect();

    // try binary name as util name.
    let args0 = args[0].clone();
    let binary = Path::new(&args0[..]);
    let binary_as_util = binary.file_name().unwrap().to_str().unwrap();

    if let Some(&uumain) = umap.get(binary_as_util) {
        std::process::exit(uumain(args));
    }

    if binary_as_util.ends_with("uutils") || binary_as_util.starts_with("uutils")
        || binary_as_util.ends_with("busybox") || binary_as_util.starts_with("busybox")
    {
        args.remove(0);
    } else {
        let mut found = false;
        for util in umap.keys() {
            if binary_as_util.ends_with(util) {
                args[0] = (*util).to_owned();
                found = true;
                break;
            }
        }
        if !found {
            println!("{}: applet not found", binary_as_util);
            std::process::exit(1);
        }
    }

    // try first arg as util name.
    if args.len() >= 1 {
        let util = &args[0][..];

        match umap.get(util) {
            Some(&uumain) => {
                std::process::exit(uumain(args.clone()));
            }
            None => {
                if &args[0][..] == "--help" || &args[0][..] == "-h" {
                    // see if they want help on a specific util
                    if args.len() >= 2 {
                        let util = &args[1][..];
                        match umap.get(util) {
                            Some(&uumain) => {
                                let code = uumain(vec![util.to_owned(), "--help".to_owned()]);
                                std::io::stdout().flush().expect("could not flush stdout");
                                std::process::exit(code);
                            }
                            None => {
                                println!("{}: applet not found", util);
                                std::process::exit(1);
                            }
                        }
                    }
                    usage(&umap);
                    std::process::exit(0);
                } else {
                    println!("{}: applet not found", util);
                    std::process::exit(1);
                }
            }
        }
    } else {
        // no arguments provided
        usage(&umap);
        std::process::exit(0);
    }
}
<commit_msg>uutils: Replace file_name with file_stem in symlink applet detection.<commit_after>#![crate_name = "uutils"]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Michael Gehring <mg@ebfe.org>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

include!(concat!(env!("OUT_DIR"), "/uutils_crates.rs"));

use std::collections::hash_map::HashMap;
use std::path::Path;
use std::env;
use std::io::Write;

extern crate uucore;

static NAME: &'static str = "uutils";
static VERSION: &'static str = env!("CARGO_PKG_VERSION");

include!(concat!(env!("OUT_DIR"), "/uutils_map.rs"));

fn usage(cmap: &UtilityMap) {
    println!("{} {}", NAME, VERSION);
    println!("");
    println!("Usage:");
    println!("  {} [util [arguments...]]\n", NAME);
    println!("Currently defined functions:");
    let mut utils: Vec<&str> = cmap.keys().map(|&s| s).collect();
    utils.sort();
    for util in utils {
        println!("\t{}", util);
    }
}

fn main() {
    uucore::panic::install_sigpipe_hook();

    let umap = util_map();
    let mut args: Vec<String> = env::args().collect();

    // try binary name as util name.
    let args0 = args[0].clone();
    let binary = Path::new(&args0[..]);
    let binary_as_util = binary.file_stem().unwrap().to_str().unwrap();

    if let Some(&uumain) = umap.get(binary_as_util) {
        std::process::exit(uumain(args));
    }

    if binary_as_util.ends_with("uutils") || binary_as_util.starts_with("uutils")
        || binary_as_util.ends_with("busybox") || binary_as_util.starts_with("busybox")
    {
        args.remove(0);
    } else {
        let mut found = false;
        for util in umap.keys() {
            if binary_as_util.ends_with(util) {
                args[0] = (*util).to_owned();
                found = true;
                break;
            }
        }
        if !found {
            println!("{}: applet not found", binary_as_util);
            std::process::exit(1);
        }
    }

    // try first arg as util name.
    if args.len() >= 1 {
        let util = &args[0][..];

        match umap.get(util) {
            Some(&uumain) => {
                std::process::exit(uumain(args.clone()));
            }
            None => {
                if &args[0][..] == "--help" || &args[0][..] == "-h" {
                    // see if they want help on a specific util
                    if args.len() >= 2 {
                        let util = &args[1][..];
                        match umap.get(util) {
                            Some(&uumain) => {
                                let code = uumain(vec![util.to_owned(), "--help".to_owned()]);
                                std::io::stdout().flush().expect("could not flush stdout");
                                std::process::exit(code);
                            }
                            None => {
                                println!("{}: applet not found", util);
                                std::process::exit(1);
                            }
                        }
                    }
                    usage(&umap);
                    std::process::exit(0);
                } else {
                    println!("{}: applet not found", util);
                    std::process::exit(1);
                }
            }
        }
    } else {
        // no arguments provided
        usage(&umap);
        std::process::exit(0);
    }
}
<|endoftext|>


<commit_before>//! Handle the decrypt TCP async via rust TLS and expose HTTP
//! TCP <=> TLS <-> HTTP
//! <=> is async
//! <-> is sync

use std::sync::Arc;
use std::pin::Pin;

use std::io::{Read as SyncRead, Write as SyncWrite};
use std::vec::Vec;

use async_std::prelude::*;
use async_std::io::{Read, Result as IoResult, Write};
use async_std::net::TcpStream;
use async_std::task::{Context, Poll};
use rustls::{ClientConfig, ClientSession, ProtocolVersion, Session};
use webpki::DNSNameRef;
use webpki_roots;

use super::errors::CabotError;
use super::results::CabotResult;

const BUFFER_PAGE_SIZE: usize = 4096;
// const RESPONSE_BUFFER_SIZE: usize = 1024;

fn create_config() -> Arc<ClientConfig> {
    let mut config = ClientConfig::new();
    config
        .root_store
        .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
    let rc_config = Arc::new(config);
    rc_config
}

fn create_client(host: &str) -> CabotResult<ClientSession> {
    let rc_config = create_config();
    let host = DNSNameRef::try_from_ascii_str(host)
        .map_err(|_| CabotError::HostnameParseError(host.to_string()))?;
    let tlsclient = ClientSession::new(&rc_config, host);
    Ok(tlsclient)
}

// fn read_buf<T>(client: &mut T, buf: &mut [u8]) -> Vec<u8>
// where
//     T: SyncRead + Sized,
// {
//     let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
//     while let Ok(count) = client.read(&mut buf[..]) {
//         if count > 0 {
//             response.extend_from_slice(&buf[0..count]);
//         } else {
//             break;
//         }
//     }
//     response
// }

pub struct TLSStream<'a> {
    tcpstream: &'a mut TcpStream,
    buf_tlsread: Vec<u8>,
    buf_tlswrite: Vec<u8>,
    buf_response: Vec<u8>,
    tlsclient: ClientSession,
}

impl<'a> TLSStream<'a> {
    pub fn new(tcpstream: &'a mut TcpStream, host: &str) -> CabotResult<Self> {
        Ok(TLSStream {
            tcpstream,
            buf_tlsread: Vec::new(),
            buf_tlswrite: Vec::new(),
            buf_response: Vec::new(),
            tlsclient: create_client(host)?,
        })
    }
    pub async fn handshake(&mut self) -> CabotResult<()> {
        let mut is_handshaking = true;
        while is_handshaking {
            while self.tlsclient.wants_write() {
                let count = self.tlsclient.write_tls(&mut self.buf_tlswrite).unwrap();
                debug!("Write {} TLS bytes", count);
                //debug!("{}", self.buf_tlswrite.len());
                self.tcpstream.write_all(&self.buf_tlswrite.as_slice()[0 .. count]).await?;
                self.tcpstream.flush().await?;
                self.buf_tlswrite.clear();
            }
            if self.tlsclient.wants_read() {
                let mut buf: [u8; BUFFER_PAGE_SIZE] = [0; BUFFER_PAGE_SIZE];
                let n = self.tcpstream.read(&mut buf[..]).await?;
                if n > 0 {
                    debug!("Read {} TCP bytes", n);
                    self.buf_tlsread.extend_from_slice(&buf[0 .. n]);
                    let count = self.tlsclient.read_tls(&mut self.buf_tlsread.as_slice())?;
                    debug!("Decode {} TLS bytes", count);

                    self.tlsclient.process_new_packets()?;

                    //let mut part: Vec<u8> = read_buf(&mut self.tlsclient, &mut self.buf_tlsread);
                    //self.buf_response.append(&mut part);
                    self.buf_tlsread.clear();
                    //self.buf_response.clear();
                }
            }
            if is_handshaking && !self.tlsclient.is_handshaking() {
                info!("Handshake complete");
                is_handshaking = false;
                let protocol = self.tlsclient.get_protocol_version();
                match protocol {
                    Some(ProtocolVersion::SSLv2) => {
                        info!("Protocol SSL v2 negociated");
                    }
                    Some(ProtocolVersion::SSLv3) => {
                        info!("Protocol SSL v3 negociated");
                    }
                    Some(ProtocolVersion::TLSv1_0) => {
                        info!("Protocol TLS v1.0 negociated");
                    }
                    Some(ProtocolVersion::TLSv1_1) => {
                        info!("Protocol TLS v1.1 negociated");
                    }
                    Some(ProtocolVersion::TLSv1_2) => {
                        info!("Protocol TLS v1.2 negociated");
                    }
                    Some(ProtocolVersion::TLSv1_3) => {
                        info!("Protocol TLS v1.3 negociated");
                    }
                    Some(ProtocolVersion::Unknown(num)) => {
                        info!("Unknown TLS Protocol negociated: {}", num);
                    }
                    None => {
                        info!("No TLS Protocol negociated");
                    }
                }
            }
        }
        Ok(())
    }

}

impl<'a> Read for TLSStream<'a> {

    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<IoResult<usize>> {

        let self_ = Pin::get_mut(self);

        let mut read_len = 0;
        let mut read_len_tls: usize = 0;

        let mut buf_tlsread = [0; BUFFER_PAGE_SIZE];
        let n = futures_core::ready!(
            Pin::new(&mut self_.tcpstream).poll_read(cx, &mut buf_tlsread)
        );

        match n {
            Err(err) => return Poll::Ready(Err(err)),
            Ok(count) => {
                let slice = &buf_tlsread[..count];
                self_.buf_tlsread.extend_from_slice(&slice);
            }
        }

        let mut bufread = self_.buf_tlsread.as_slice();

        while self_.tlsclient.wants_read() {
            let count = self_.tlsclient.read_tls(&mut bufread)?;
            read_len_tls = read_len_tls + count;
            debug!("Read {} TLS bytes", count);
            if count == 0 {
                debug!("Break");
                break;
            }
        }

        self_.tlsclient.process_new_packets().unwrap();

        let mut bufr = [0; BUFFER_PAGE_SIZE];
        let ret = self_.tlsclient.read(&mut bufr[..]);
        if let Ok(count) = ret {
            read_len = read_len + count;
            if count > 0 {
                self_.buf_response.extend_from_slice(&buf[0..count]);
            }
        }
        Poll::Ready(Ok(read_len))
    }
}

impl<'a> Write for TLSStream<'a> {

    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<IoResult<usize>> {
        let self_ = Pin::get_mut(self);

        self_.tlsclient.write_all(&buf).unwrap();
        let stream = Pin::new(&mut *self_.tcpstream);

        while self_.tlsclient.wants_write() {
            let count = self_.tlsclient.write_tls(&mut self_.buf_tlswrite).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        let n = futures_core::ready!(
            stream.poll_write(cx, self_.buf_tlswrite.as_slice())
        )?;
        self_.buf_tlswrite.clear();
        Poll::Ready(Ok(n))
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<IoResult<()>> {
        let self_ = Pin::get_mut(self);
        Pin::new(&mut *self_.tcpstream).poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<IoResult<()>> {
        let self_ = Pin::get_mut(self);
        Pin::new(&mut *self_.tcpstream).poll_close(cx)
    }

}

<commit_msg>Use self.read during handshake<commit_after>//! Handle the decrypt TCP async via rust TLS and expose HTTP
//! TCP <=> TLS <-> HTTP
//! <=> is async
//! <-> is sync

use std::pin::Pin;
use std::sync::Arc;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Read as SyncRead, Write as SyncWrite};
use std::vec::Vec;

use async_std::io::{Read, Result as IoResult, Write};
use async_std::net::TcpStream;
use async_std::prelude::*;
use async_std::task::{Context, Poll};
use rustls::{ClientConfig, ClientSession, ProtocolVersion, Session};
use webpki::DNSNameRef;
use webpki_roots;

use super::errors::CabotError;
use super::results::CabotResult;

const BUFFER_PAGE_SIZE: usize = 4096;
// const RESPONSE_BUFFER_SIZE: usize = 1024;

fn create_config() -> Arc<ClientConfig> {
    let mut config = ClientConfig::new();
    config
        .root_store
        .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
    let rc_config = Arc::new(config);
    rc_config
}

fn create_client(host: &str) -> CabotResult<ClientSession> {
    let rc_config = create_config();
    let host = DNSNameRef::try_from_ascii_str(host)
        .map_err(|_| CabotError::HostnameParseError(host.to_string()))?;
    let tlsclient = ClientSession::new(&rc_config, host);
    Ok(tlsclient)
}

pub struct TLSStream<'a> {
    tcpstream: &'a mut TcpStream,
    buf_tlsread: Vec<u8>,
    buf_tlswrite: Vec<u8>,
    tlsclient: ClientSession,
}

impl<'a> TLSStream<'a> {
    pub fn new(tcpstream: &'a mut TcpStream, host: &str) -> CabotResult<Self> {
        Ok(TLSStream {
            tcpstream,
            buf_tlsread: Vec::new(),
            buf_tlswrite: Vec::new(),
            tlsclient: create_client(host)?,
        })
    }
    pub async fn handshake(&mut self) -> CabotResult<()> {
        let mut is_handshaking = true;
        let mut read_buf: [u8; BUFFER_PAGE_SIZE] = [0; BUFFER_PAGE_SIZE];
        while is_handshaking {
            while self.tlsclient.wants_write() {
                let count = self.tlsclient.write_tls(&mut self.buf_tlswrite).unwrap();
                debug!("Write {} TLS bytes", count);
                //debug!("{}", self.buf_tlswrite.len());
                self.tcpstream
                    .write_all(&self.buf_tlswrite.as_slice()[0..count])
                    .await?;
                self.tcpstream.flush().await?;
                self.buf_tlswrite.clear();
            }
            if self.tlsclient.wants_read() {
                let count = self.read(&mut read_buf).await?;
                debug!("Read {} TLS bytes during hangshake", count);
            }
            if is_handshaking && !self.tlsclient.is_handshaking() {
                info!("Handshake complete");
                is_handshaking = false;
                let protocol = self.tlsclient.get_protocol_version();
                match protocol {
                    Some(ProtocolVersion::SSLv2) => {
                        info!("Protocol SSL v2 negociated");
                    }
                    Some(ProtocolVersion::SSLv3) => {
                        info!("Protocol SSL v3 negociated");
                    }
                    Some(ProtocolVersion::TLSv1_0) => {
                        info!("Protocol TLS v1.0 negociated");
                    }
                    Some(ProtocolVersion::TLSv1_1) => {
                        info!("Protocol TLS v1.1 negociated");
                    }
                    Some(ProtocolVersion::TLSv1_2) => {
                        info!("Protocol TLS v1.2 negociated");
                    }
                    Some(ProtocolVersion::TLSv1_3) => {
                        info!("Protocol TLS v1.3 negociated");
                    }
                    Some(ProtocolVersion::Unknown(num)) => {
                        info!("Unknown TLS Protocol negociated: {}", num);
                    }
                    None => {
                        info!("No TLS Protocol negociated");
                    }
                }
            }
        }
        Ok(())
    }
}

impl<'a> Read for TLSStream<'a> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<IoResult<usize>> {
        let self_ = Pin::get_mut(self);
        let mut tcp_buf: [u8; BUFFER_PAGE_SIZE] = [0; BUFFER_PAGE_SIZE];

        let count =
            futures_core::ready!(Pin::new(&mut self_.tcpstream).poll_read(cx, &mut tcp_buf[..]));

        match count {
            Err(err) => Poll::Ready(Err(err)),
            Ok(n) => {
                if n > 0 {
                    debug!("Read {} TCP bytes", n);
                    self_.buf_tlsread.extend_from_slice(&tcp_buf[0..n]);
                    let count = self_
                        .tlsclient
                        .read_tls(&mut self_.buf_tlsread.as_slice())?;
                    debug!("Decode {} TLS bytes", count);

                    let packets = self_.tlsclient.process_new_packets();
                    match packets {
                        Ok(_) => {
                            let cnt = self_.tlsclient.read(&mut buf[..]).unwrap();
                            debug!("Read {} Unencrypted bytes", cnt);
                            let _: Vec<_> = self_.buf_tlsread.drain(..cnt).collect();
                            //debug!("??? {} ???", self_.buf_tlsread.len());
                            //debug!("??? {} ???", n.len());
                            self_.buf_tlsread.clear();
                            Poll::Ready(Ok(cnt))
                        }
                        Err(err) => Poll::Ready(Err(IoError::new(
                            IoErrorKind::InvalidData,
                            format!("{:?}", err),
                        ))),
                    }
                } else {
                    Poll::Pending
                }
            }
        }
    }
}

impl<'a> Write for TLSStream<'a> {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<IoResult<usize>> {
        let self_ = Pin::get_mut(self);

        self_.tlsclient.write_all(&buf).unwrap();
        let stream = Pin::new(&mut *self_.tcpstream);

        while self_.tlsclient.wants_write() {
            let count = self_.tlsclient.write_tls(&mut self_.buf_tlswrite).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        let n = futures_core::ready!(stream.poll_write(cx, self_.buf_tlswrite.as_slice()))?;
        self_.buf_tlswrite.clear();
        Poll::Ready(Ok(n))
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<IoResult<()>> {
        let self_ = Pin::get_mut(self);
        Pin::new(&mut *self_.tcpstream).poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<IoResult<()>> {
        let self_ = Pin::get_mut(self);
        Pin::new(&mut *self_.tcpstream).poll_close(cx)
    }
}
<|endoftext|>


<commit_before>use std::time::Duration;

use disque::Disque;
use redis::{RedisResult, Iter};

/// Helper to add a new job
///
/// # Examples
///
/// ```
/// # use disque::Disque;
/// # use disque::AddJobBuilder;
///
/// let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
/// let jobid = AddJobBuilder::new(b"example queue", b"my job", 10000
///     ).delay(1440).run(&disque).unwrap();
/// ```
pub struct AddJobBuilder<'a> {
    queue_name: &'a [u8],
    job: &'a [u8],
    timeout: Duration,
    replicate: Option<usize>,
    delay: Option<Duration>,
    retry: Option<Duration>,
    ttl: Option<Duration>,
    maxlen: Option<usize>,
    async: bool,
}

impl<'a> AddJobBuilder<'a> {
    /// Creates a new builder for `queue_name`.
    /// Timeout is specified in milliseconds.
    pub fn new(queue_name: &'a [u8], job: &'a [u8], timeout_ms: u64
            ) -> AddJobBuilder<'a> {
        AddJobBuilder {
            queue_name: queue_name,
            job: job,
            timeout: Duration::from_millis(timeout_ms),
            replicate: None,
            delay: None,
            retry: None,
            ttl: None,
            maxlen: None,
            async: false,
        }
    }

    /// Changes the queue name where the job will be added.
    pub fn queue_name(&mut self, queue_name: &'a [u8]) -> &mut Self {
        self.queue_name = queue_name; self
    }

    /// Changes the job body.
    pub fn job(&mut self, job: &'a [u8]) -> &mut Self {
        self.job = job; self
    }

    /// Changes the timeout. It must be specified in milliseconds.
    pub fn timeout(&mut self, timeout_ms: u64) -> &mut Self {
        self.timeout = Duration::from_millis(timeout_ms); self
    }

    /// The number of nodes the job should be replicated to.
    pub fn replicate(&mut self, replicate: usize) -> &mut Self {
        self.replicate = Some(replicate); self
    }

    /// The number of seconds that should elapse before the job is queued.
    pub fn delay(&mut self, delay: u64) -> &mut Self {
        self.delay = Some(Duration::from_secs(delay)); self
    }

    /// Period after which, if no ACK is received, the job is put again
    /// into the queue for delivery
    pub fn retry(&mut self, retry: u64) -> &mut Self {
        self.retry = Some(Duration::from_secs(retry)); self
    }

    /// The max job life in seconds. After this time, the job is deleted even
    /// if it was not successfully delivered.
    pub fn ttl(&mut self, ttl: u64) -> &mut Self {
        self.ttl = Some(Duration::from_secs(ttl)); self
    }

    /// If there are already count messages queued for the specified queue
    /// name, the message is refused and an error reported to the client.
    pub fn maxlen(&mut self, maxlen: usize) -> &mut Self {
        self.maxlen = Some(maxlen); self
    }

    /// If true, asks the server to let the command return ASAP and replicate
    /// the job to other nodes in the background.
    /// Otherwise, the job is put into the queue only when the client gets a
    /// positive reply.
    pub fn async(&mut self, async: bool) -> &mut Self {
        self.async = async; self
    }

    /// Executes the addjob command.
    pub fn run(&self, disque: &Disque) -> RedisResult<String> {
        disque.addjob(self.queue_name, self.job, self.timeout, self.replicate,
                self.delay, self.retry, self.ttl, self.maxlen, self.async)
    }
}

#[test]
fn add_job_builder() {
    let mut jb = AddJobBuilder::new(b"queue", b"job", 123);
    assert_eq!(jb.queue_name, b"queue");
    assert_eq!(jb.job, b"job");
    assert_eq!(jb.timeout, Duration::from_millis(123));
    assert_eq!(jb.replicate, None);
    assert_eq!(jb.delay, None);
    assert_eq!(jb.retry, None);
    assert_eq!(jb.ttl, None);
    assert_eq!(jb.maxlen, None);
    assert_eq!(jb.async, false);
    jb.replicate(3).delay(4).retry(5).ttl(6).maxlen(7).async(true);
    assert_eq!(jb.queue_name, b"queue");
    assert_eq!(jb.job, b"job");
    assert_eq!(jb.timeout, Duration::from_millis(123));
    assert_eq!(jb.replicate, Some(3));
    assert_eq!(jb.delay, Some(Duration::from_secs(4)));
    assert_eq!(jb.retry, Some(Duration::from_secs(5)));
    assert_eq!(jb.ttl, Some(Duration::from_secs(6)));
    assert_eq!(jb.maxlen, Some(7));
    assert_eq!(jb.async, true);
    jb.queue_name(b"my queue").job(b"my job").timeout(234);
    assert_eq!(jb.queue_name, b"my queue");
    assert_eq!(jb.job, b"my job");
    assert_eq!(jb.timeout, Duration::from_millis(234));
    assert_eq!(jb.replicate, Some(3));
    assert_eq!(jb.delay, Some(Duration::from_secs(4)));
    assert_eq!(jb.retry, Some(Duration::from_secs(5)));
    assert_eq!(jb.ttl, Some(Duration::from_secs(6)));
    assert_eq!(jb.maxlen, Some(7));
    assert_eq!(jb.async, true);
}

/// Helper to get a list of queues
///
/// # Examples
///
/// ```
/// # use disque::Disque;
/// # use disque::QueueQueryBuilder;
///
/// let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
/// let queues = QueueQueryBuilder::new().busyloop(true).minlen(50)
///     .iter(&disque).unwrap().collect::<Vec<_>>();
/// assert!(queues.len() >= 0);
/// ```
pub struct QueueQueryBuilder {
    count: u64,
    busyloop: bool,
    minlen: Option<u64>,
    maxlen: Option<u64>,
    importrate: Option<u64>,
}

impl QueueQueryBuilder {
    /// Creates a new builder.
    pub fn new() -> QueueQueryBuilder {
        QueueQueryBuilder {
            count: 16,
            busyloop: false,
            minlen: None,
            maxlen: None,
            importrate: None,
        }
    }

    /// A hint about how much work to do per iteration.
    pub fn count(&mut self, count: u64) -> &mut Self {
        self.count = count; self
    }

    /// If true, blocks and returns all the queues in a busy loop.
    pub fn busyloop(&mut self, busyloop: bool) -> &mut Self {
        self.busyloop = busyloop; self
    }

    /// Only return queues with at least `minlen` jobs.
    pub fn minlen(&mut self, minlen: u64) -> &mut Self {
        self.minlen = Some(minlen); self
    }

    /// Only return queues with at most `maxlen` jobs.
    pub fn maxlen(&mut self, maxlen: u64) -> &mut Self {
        self.maxlen = Some(maxlen); self
    }

    /// Only return queues with a job import rate (from other nodes) greater
    /// than or equal to  `importrate`.
    pub fn importrate(&mut self, importrate: u64) -> &mut Self {
        self.importrate = Some(importrate); self
    }

    /// Gets the queue iterator.
    pub fn iter<'a>(&'a self, disque: &'a Disque) -> RedisResult<Iter<Vec<u8>>> {
        disque.qscan(0, self.count, self.busyloop, self.minlen, self.maxlen,
                self.importrate)
    }
}

#[test]
fn queue_query_builder() {
    let mut qqb = QueueQueryBuilder::new();
    assert_eq!(qqb.count, 16);
    assert_eq!(qqb.busyloop, false);
    assert_eq!(qqb.minlen, None);
    assert_eq!(qqb.maxlen, None);
    assert_eq!(qqb.importrate, None);
    qqb.count(11).busyloop(true).minlen(1).maxlen(10).importrate(5);
    assert_eq!(qqb.count, 11);
    assert_eq!(qqb.busyloop, true);
    assert_eq!(qqb.minlen, Some(1));
    assert_eq!(qqb.maxlen, Some(10));
    assert_eq!(qqb.importrate, Some(5));
}
<commit_msg>Add job query builder<commit_after>use std::collections::HashMap;
use std::time::Duration;

use disque::Disque;
use redis::{RedisResult, Iter, Value};

/// Helper to add a new job
///
/// # Examples
///
/// ```
/// # use disque::Disque;
/// # use disque::AddJobBuilder;
///
/// let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
/// let jobid = AddJobBuilder::new(b"example queue", b"my job", 10000
///     ).delay(1440).run(&disque).unwrap();
/// ```
pub struct AddJobBuilder<'a> {
    queue_name: &'a [u8],
    job: &'a [u8],
    timeout: Duration,
    replicate: Option<usize>,
    delay: Option<Duration>,
    retry: Option<Duration>,
    ttl: Option<Duration>,
    maxlen: Option<usize>,
    async: bool,
}

impl<'a> AddJobBuilder<'a> {
    /// Creates a new builder for `queue_name`.
    /// Timeout is specified in milliseconds.
    pub fn new(queue_name: &'a [u8], job: &'a [u8], timeout_ms: u64
            ) -> AddJobBuilder<'a> {
        AddJobBuilder {
            queue_name: queue_name,
            job: job,
            timeout: Duration::from_millis(timeout_ms),
            replicate: None,
            delay: None,
            retry: None,
            ttl: None,
            maxlen: None,
            async: false,
        }
    }

    /// Changes the queue name where the job will be added.
    pub fn queue_name(&mut self, queue_name: &'a [u8]) -> &mut Self {
        self.queue_name = queue_name; self
    }

    /// Changes the job body.
    pub fn job(&mut self, job: &'a [u8]) -> &mut Self {
        self.job = job; self
    }

    /// Changes the timeout. It must be specified in milliseconds.
    pub fn timeout(&mut self, timeout_ms: u64) -> &mut Self {
        self.timeout = Duration::from_millis(timeout_ms); self
    }

    /// The number of nodes the job should be replicated to.
    pub fn replicate(&mut self, replicate: usize) -> &mut Self {
        self.replicate = Some(replicate); self
    }

    /// The number of seconds that should elapse before the job is queued.
    pub fn delay(&mut self, delay: u64) -> &mut Self {
        self.delay = Some(Duration::from_secs(delay)); self
    }

    /// Period after which, if no ACK is received, the job is put again
    /// into the queue for delivery
    pub fn retry(&mut self, retry: u64) -> &mut Self {
        self.retry = Some(Duration::from_secs(retry)); self
    }

    /// The max job life in seconds. After this time, the job is deleted even
    /// if it was not successfully delivered.
    pub fn ttl(&mut self, ttl: u64) -> &mut Self {
        self.ttl = Some(Duration::from_secs(ttl)); self
    }

    /// If there are already count messages queued for the specified queue
    /// name, the message is refused and an error reported to the client.
    pub fn maxlen(&mut self, maxlen: usize) -> &mut Self {
        self.maxlen = Some(maxlen); self
    }

    /// If true, asks the server to let the command return ASAP and replicate
    /// the job to other nodes in the background.
    /// Otherwise, the job is put into the queue only when the client gets a
    /// positive reply.
    pub fn async(&mut self, async: bool) -> &mut Self {
        self.async = async; self
    }

    /// Executes the addjob command.
    pub fn run(&self, disque: &Disque) -> RedisResult<String> {
        disque.addjob(self.queue_name, self.job, self.timeout, self.replicate,
                self.delay, self.retry, self.ttl, self.maxlen, self.async)
    }
}

#[test]
fn add_job_builder() {
    let mut jb = AddJobBuilder::new(b"queue", b"job", 123);
    assert_eq!(jb.queue_name, b"queue");
    assert_eq!(jb.job, b"job");
    assert_eq!(jb.timeout, Duration::from_millis(123));
    assert_eq!(jb.replicate, None);
    assert_eq!(jb.delay, None);
    assert_eq!(jb.retry, None);
    assert_eq!(jb.ttl, None);
    assert_eq!(jb.maxlen, None);
    assert_eq!(jb.async, false);
    jb.replicate(3).delay(4).retry(5).ttl(6).maxlen(7).async(true);
    assert_eq!(jb.queue_name, b"queue");
    assert_eq!(jb.job, b"job");
    assert_eq!(jb.timeout, Duration::from_millis(123));
    assert_eq!(jb.replicate, Some(3));
    assert_eq!(jb.delay, Some(Duration::from_secs(4)));
    assert_eq!(jb.retry, Some(Duration::from_secs(5)));
    assert_eq!(jb.ttl, Some(Duration::from_secs(6)));
    assert_eq!(jb.maxlen, Some(7));
    assert_eq!(jb.async, true);
    jb.queue_name(b"my queue").job(b"my job").timeout(234);
    assert_eq!(jb.queue_name, b"my queue");
    assert_eq!(jb.job, b"my job");
    assert_eq!(jb.timeout, Duration::from_millis(234));
    assert_eq!(jb.replicate, Some(3));
    assert_eq!(jb.delay, Some(Duration::from_secs(4)));
    assert_eq!(jb.retry, Some(Duration::from_secs(5)));
    assert_eq!(jb.ttl, Some(Duration::from_secs(6)));
    assert_eq!(jb.maxlen, Some(7));
    assert_eq!(jb.async, true);
}

/// Helper to get a list of queues
///
/// # Examples
///
/// ```
/// # use disque::Disque;
/// # use disque::QueueQueryBuilder;
///
/// let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
/// let queues = QueueQueryBuilder::new().busyloop(true).minlen(50)
///     .iter(&disque).unwrap().collect::<Vec<_>>();
/// assert!(queues.len() >= 0);
/// ```
pub struct QueueQueryBuilder {
    count: u64,
    busyloop: bool,
    minlen: Option<u64>,
    maxlen: Option<u64>,
    importrate: Option<u64>,
}

impl QueueQueryBuilder {
    /// Creates a new builder.
    pub fn new() -> QueueQueryBuilder {
        QueueQueryBuilder {
            count: 16,
            busyloop: false,
            minlen: None,
            maxlen: None,
            importrate: None,
        }
    }

    /// A hint about how much work to do per iteration.
    pub fn count(&mut self, count: u64) -> &mut Self {
        self.count = count; self
    }

    /// If true, blocks and returns all the queues in a busy loop.
    pub fn busyloop(&mut self, busyloop: bool) -> &mut Self {
        self.busyloop = busyloop; self
    }

    /// Only return queues with at least `minlen` jobs.
    pub fn minlen(&mut self, minlen: u64) -> &mut Self {
        self.minlen = Some(minlen); self
    }

    /// Only return queues with at most `maxlen` jobs.
    pub fn maxlen(&mut self, maxlen: u64) -> &mut Self {
        self.maxlen = Some(maxlen); self
    }

    /// Only return queues with a job import rate (from other nodes) greater
    /// than or equal to  `importrate`.
    pub fn importrate(&mut self, importrate: u64) -> &mut Self {
        self.importrate = Some(importrate); self
    }

    /// Gets the queue iterator.
    pub fn iter<'a>(&'a self, disque: &'a Disque) -> RedisResult<Iter<Vec<u8>>> {
        disque.qscan(0, self.count, self.busyloop, self.minlen, self.maxlen,
                self.importrate)
    }
}

#[test]
fn queue_query_builder() {
    let mut qqb = QueueQueryBuilder::new();
    assert_eq!(qqb.count, 16);
    assert_eq!(qqb.busyloop, false);
    assert_eq!(qqb.minlen, None);
    assert_eq!(qqb.maxlen, None);
    assert_eq!(qqb.importrate, None);
    qqb.count(11).busyloop(true).minlen(1).maxlen(10).importrate(5);
    assert_eq!(qqb.count, 11);
    assert_eq!(qqb.busyloop, true);
    assert_eq!(qqb.minlen, Some(1));
    assert_eq!(qqb.maxlen, Some(10));
    assert_eq!(qqb.importrate, Some(5));
}

/// Helper to get a list of jobs
///
/// # Examples
///
/// ```
/// # use disque::Disque;
/// # use disque::JobQueryBuilder;
///
/// let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
/// let jobs = JobQueryBuilder::new().queue(b"my queue").state("queued")
///     .iter_ids(&disque).unwrap().collect::<Vec<_>>();
/// assert!(jobs.len() >= 0);
/// ```
pub struct JobQueryBuilder<'a> {
    count: u64,
    blocking: bool,
    queue: Option<&'a [u8]>,
    states: Vec<&'a str>,
}

impl<'a> JobQueryBuilder<'a> {
    pub fn new() -> JobQueryBuilder<'a> {
        JobQueryBuilder {
            count: 16,
            blocking: false,
            queue: None,
            states: vec![],
        }
    }

    pub fn count(&mut self, count: u64) -> &mut Self {
        self.count = count; self
    }

    pub fn blocking(&mut self, blocking: bool) -> &mut Self {
        self.blocking = blocking; self
    }

    pub fn queue(&mut self, queue: &'a [u8]) -> &mut Self {
        self.queue = Some(queue); self
    }

    pub fn state(&mut self, state: &'a str) -> &mut Self {
        self.states.push(state); self
    }

    /// Gets a job ids iterator.
    pub fn iter_ids<'b>(&'b self, disque: &'b Disque
            ) -> RedisResult<Iter<String>> {
        disque.jscan_id(0, self.count, self.blocking, self.queue,
                &*self.states)
    }

    /// Gets a job information iterator.
    pub fn iter_all<'b>(&'b self, disque: &'b Disque
            ) -> RedisResult<Iter<HashMap<String, Value>>> {
        disque.jscan_all(0, self.count, self.blocking, self.queue,
                &*self.states)
    }
}

#[test]
fn job_query_builder() {
    let mut jqb = JobQueryBuilder::new();
    assert_eq!(jqb.count, 16);
    assert_eq!(jqb.blocking, false);
    assert_eq!(jqb.queue, None);
    assert_eq!(jqb.states.len(), 0);
    jqb.count(20).blocking(true).queue(b"jqb queue").state("state1"
            ).state("state2");
    assert_eq!(jqb.count, 20);
    assert_eq!(jqb.blocking, true);
    assert_eq!(jqb.queue, Some(b"jqb queue" as &[u8]));
    assert_eq!(jqb.states, vec!["state1", "state2"]);
}
<|endoftext|>


<commit_before>use std::fmt::{self, Display, Debug};
use std::cell::Cell;
use std::io::Read;
use std::io;
use nom;

const DEFAULT_BUF_SIZE: usize = 8 * 1024;

pub trait StreamBuf<O> {
    fn fill(&mut self, min_count: usize) -> Result<(), io::Error>;
    fn recycle(&mut self);
    fn consume(&mut self, count: usize);
    fn data<'b>(&'b self) -> &'b[O];
    fn needed(&self) -> Option<nom::Needed>;
    fn apply<'b, C, CO, E>(&'b self, combinator: C) -> Result<CO, ApplyError<&'b [O], E>> where
        O: 'b, C: Fn(&'b [O]) -> nom::IResult<&'b [O], CO, E>;

    fn fill_apply<'b, C, CO, E>(&'b mut self, combinator: C) -> Result<Option<CO>, nom::Err<&'b [O], E>> where
        C: Fn(&'b [O]) -> nom::IResult<&'b [O], CO, E> {
        use stream_buf::ApplyError;
        use nom::{Needed};

        //TODO: io::Error
        match self.needed() {
            Some(Needed::Size(bytes)) => self.fill(bytes).expect("IO Error"),
            Some(Needed::Unknown) => self.fill(1).expect("IO Error"),
            None => ()
        }

        match self.apply(combinator) {
            Ok(out) => Ok(Some(out)),
            Err(ApplyError::Parser(err)) => Err(err),
            Err(ApplyError::TryAgain) => Ok(None),
        }
    }
}

// Note: Need to use macro as this cannot be represented in the type system.
//       It would require generic combinator outpt parameter to have for<'r> life time
// TODO: Makie it smarter and able to handle errors and unknown amount of needed data
#[macro_export]
macro_rules! apply_stream (
    ($sb:expr, $comb:expr) => ({
        use stream_buf::{StreamBuf, ApplyError};
        use nom::{Needed};

        //TODO: io::Error
        match $sb.needed() {
            Some(Needed::Size(bytes)) => $sb.fill(bytes).expect("IO Error"),
            Some(Needed::Unknown) => $sb.fill(1).expect("IO Error"),
            None => ()
        }

        match $sb.apply($comb) {
            Ok(out) => Ok(Some(out)),
            Err(ApplyError::Parser(err)) => Err(err),
            Err(ApplyError::TryAgain) => Ok(None),
        }
    })
);

/*
#[macro_export]
macro_rules! fill_and_apply_stream (
    ($sb:expr, $comb:expr) => ({
        let result;
        loop {
            match apply_stream!($sb, $comb) {
                Ok(Some(out)) => {
                    result = out;
                    break
                }
                Err(err) => {
                    //result = Err(err);
                    //break
                }
                _ => ()
            }
        }
        result
    })
);
*/

#[derive(Debug)]
pub enum ApplyError<I, E> {
    //Io(io::Error),
    Parser(nom::Err<I, E>),
    TryAgain
}

/*
impl<I, E> From<io::Error> for ApplyError<I, E> {
    fn from(e: io::Error) -> ApplyError<I, E> {
        ApplyError::Io(e)
    }
}
*/
/*
impl<I, E> From<nom::Err<I, E>> for ApplyError<I, E> {
    fn from(e: nom::Err<I, E>) -> ApplyError<I, E> {
        ApplyError::Parser(e)
    }
}

impl<I, E> From<nom::Needed> for ApplyError<I, E> {
    fn from(e: nom::Needed) -> ApplyError<I, E> {
        ApplyError::NeedMore(e)
    }
}
*/
/*
trait ToApplyError<I, O, E> {
    fn into_vsl_result(self) -> Result<O, ApplyError<I, E>>;
}

impl<I, O, E> ToApplyError<I, O, E> for nom::IResult<I, O, E> {
    fn into_vsl_result(self) -> Result<O, ApplyError<I, E>> {
        match self {
            nom::IResult::Done(_, out) => Ok(out),
            nom::IResult::Error(e) => Err(From::from(e)),
            nom::IResult::Incomplete(n) => Err(From::from(n)),
        }
    }
}
*/

impl<I, E> Display for ApplyError<I, E> where I: Debug, E: Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            //&ApplyError::Io(ref e) => write!(f, "Failed to read VSL data: {}", e),
            &ApplyError::Parser(ref e) => write!(f, "Failed to parse data: {}", e),
            //&ApplyError::NeedMore(ref e) => write!(f, "Not enought data to finish parsing: {:?}", e),
            &ApplyError::TryAgain => write!(f, "Not enought data to finish parsing - try again later"),
        }
    }
}

pub struct ReadStreamBuf<R: Read> {
    reader: R,
    //TODO: alignment?
    buf: Vec<u8>,
    cap: usize,
    needed: Cell<Option<nom::Needed>>,
    offset: Cell<usize>,
}

impl<R: Read> ReadStreamBuf<R> {
    pub fn new(reader: R) -> ReadStreamBuf<R> {
        ReadStreamBuf::with_capacity(reader, DEFAULT_BUF_SIZE)
    }

    pub fn with_capacity(reader: R, cap: usize) -> ReadStreamBuf<R> {
        ReadStreamBuf {
            reader: reader,
            buf: Vec::with_capacity(cap),
            cap: cap,
            needed: Cell::new(Some(nom::Needed::Unknown)),
            offset: Cell::new(0),
        }
    }
}

impl<R: Read> StreamBuf<u8> for ReadStreamBuf<R> {
    fn fill(&mut self, min_bytes: usize) -> Result<(), io::Error> {
        let len = self.buf.len();
        let have = len - self.offset.get();
        let needed = min_bytes - have;

        if have >= min_bytes {
            return Ok(())
        }

        //TODO: enforce cap
        self.buf.resize(len + needed, 0);
        //println!("fill needed: {}", needed);
        //println!("buf write: {}..{} ({}); have: {} will have: {}", len, len + needed, self.buf[len..len + needed].len(), have, have + needed);
        let result = self.reader.read_exact(&mut self.buf[len..len + needed]);

        if result.is_err() {
            self.buf.resize(len, 0);
        }

        //println!("buf have: {:?}", self.data());

        //TODO: Try to read all that we have non-blocking in case we have more
        result
    }

    fn recycle(&mut self) {
        if self.offset.get() == 0 {
            return
        }
        self.buf = self.buf.split_off(self.offset.get());
        self.offset.set(0);
    }

    fn consume(&mut self, bytes: usize) {
        let len = self.buf.len();

        let consume = if self.offset.get() + bytes > len {
            len - self.offset.get()
        } else {
            bytes
        };
        self.offset.set(self.offset.get() + consume);
    }

    fn data<'b>(&'b self) -> &'b[u8] {
        &self.buf[self.offset.get()..self.buf.len()]
    }

    fn needed(&self) -> Option<nom::Needed> {
        self.needed.get().clone()
    }

    fn apply<'b, C, CO, E>(&'b self, combinator: C) -> Result<CO, ApplyError<&'b [u8], E>> where
        C: Fn(&'b [u8]) -> nom::IResult<&'b [u8], CO, E> {
        let data = self.data();
        let (left, out) = match combinator(data) {
            nom::IResult::Done(left, out) => (left, out),
            nom::IResult::Error(err) => return Err(ApplyError::Parser(err)),
            nom::IResult::Incomplete(needed) => {
                //println!("incomplete: needed: {:?}", needed);
                self.needed.set(Some(needed));
                return Err(ApplyError::TryAgain)
            }
        };
        let consumed = data.len() - left.len();
        //println!("done: consumed: {}", consumed);
        // Need Cell to modify offset after parsing is done
        self.offset.set(self.offset.get() + consumed);

        // We don't know how much we will need now
        self.needed.set(None);
        return Ok(out)
    }
}

#[cfg(test)]
mod resd_stream_buf_tests {
    use super::{StreamBuf, ApplyError};
    use super::ReadStreamBuf;
    use nom::{IResult, Needed};
    use std::io::Cursor;

    fn subject(data: Vec<u8>) -> ReadStreamBuf<Cursor<Vec<u8>>> {
        ReadStreamBuf::new(Cursor::new(data))
    }

    fn subject_with_default_data() -> ReadStreamBuf<Cursor<Vec<u8>>> {
        subject(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
    }

    #[test]
    fn reading() {
        let mut rsb = subject_with_default_data();
        assert_eq!(rsb.data(), [].as_ref());

        rsb.fill(2).unwrap();
        assert_eq!(rsb.data(), [0, 1].as_ref());

        rsb.fill(3).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2].as_ref());
    }

    #[test]
    fn recycle() {
        let mut rsb = subject_with_default_data();

        rsb.fill(3).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2].as_ref());

        rsb.recycle();
        assert_eq!(rsb.data(), [0, 1, 2].as_ref());

        rsb.fill(5).unwrap();
        rsb.fill(7).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4, 5, 6].as_ref());
    }

    #[test]
    fn consume() {
        let mut rsb = subject_with_default_data();

        rsb.fill(5).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4].as_ref());

        rsb.consume(2);
        assert_eq!(rsb.data(), [2, 3, 4].as_ref());
    }

    #[test]
    fn consume_more_than_we_have() {
        let mut rsb = subject_with_default_data();

        // We don't consume more than we have already in the buffor
        rsb.fill(2).unwrap();
        rsb.consume(3);
        rsb.fill(2).unwrap();
        assert_eq!(rsb.data(), [2, 3].as_ref());
    }

    #[test]
    fn apply_function() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), 0);
    }

    #[test]
    fn apply_should_consume() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), 0);
        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), 1);
    }

    #[test]
    fn apply_converted_macro() {
        let mut rsb = subject_with_default_data();

        rsb.fill(2).unwrap();
        assert_eq!(rsb.apply(closure!(tag!([0, 1]))).unwrap(), [0, 1].as_ref());
    }

    #[test]
    fn apply_closure() {
        let mut rsb = subject_with_default_data();

        rsb.fill(2).unwrap();
        assert_eq!(rsb.apply(|i| tag!(i, [0, 1])).unwrap(), [0, 1].as_ref());
    }

    #[test]
    fn apply_custom_fuction_with_refs() {
        let mut rsb = subject_with_default_data();
        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }

        rsb.fill(3).unwrap();
        assert_eq!(rsb.apply(comb).unwrap(), [0, 1, 2].as_ref());
    }

    #[test]
    fn needed_with_apply() {
        let mut rsb = subject_with_default_data();

        fn comb<'a>(input: &'a[u8]) -> IResult<&'a[u8], &'a[u8]> {
            tag!(input, [0, 1, 2])
        }

        if let ApplyError::TryAgain = rsb.apply(comb).unwrap_err() {
        } else {
            assert!(false) //TODO: fix
        }

        let needed = rsb.needed();
        assert_eq!(needed, Some(Needed::Size(3)));

        if let Some(Needed::Size(bytes)) = needed {
            rsb.fill(bytes).unwrap();
            assert_eq!(rsb.apply(comb).unwrap(), [0, 1, 2].as_ref());
        }
    }

    #[test]
    fn apply_stream() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }

        /*
        assert_eq!(apply_stream!(rsb, comb), Ok(None));
        assert_eq!(apply_stream!(rsb, comb), Ok(Some([0, 1, 2].as_ref())));
        assert_eq!(apply_stream!(rsb, be_u8), Ok(None));
        assert_eq!(apply_stream!(rsb, be_u8), Ok(Some(3)));
        */

        assert_eq!(rsb.fill_apply(comb), Ok(None));
        assert_eq!(rsb.fill_apply(comb), Ok(Some([0, 1, 2].as_ref())));
        assert_eq!(rsb.fill_apply(be_u8), Ok(None));
        assert_eq!(rsb.fill_apply(be_u8), Ok(Some(3)));
    }

    /*
    #[test]
    fn fill_and_apply_stream() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }
        fill_and_apply_stream!(rsb, comb);
        //assert_eq!(fill_and_apply_stream!(rsb, comb), Ok(Some([0, 1, 2].as_ref())));
        //assert_eq!(fill_and_apply_stream!(rsb, be_u8), Ok(Some(3)));
    }
    */
}
<commit_msg>cleanup<commit_after>use std::fmt::{self, Display, Debug};
use std::cell::Cell;
use std::io::Read;
use std::io;
use nom;

const DEFAULT_BUF_SIZE: usize = 8 * 1024;

pub trait StreamBuf<O> {
    fn fill(&mut self, min_count: usize) -> Result<(), io::Error>;
    fn recycle(&mut self);
    fn consume(&mut self, count: usize);
    fn data<'b>(&'b self) -> &'b[O];
    fn needed(&self) -> Option<nom::Needed>;
    fn apply<'b, C, CO, E>(&'b self, combinator: C) -> Result<Option<CO>, nom::Err<&'b [O], E>> where
        O: 'b, C: Fn(&'b [O]) -> nom::IResult<&'b [O], CO, E>;

    fn fill_apply<'b, C, CO, E>(&'b mut self, combinator: C) -> Result<Option<CO>, FillApplyError<&'b [O], E>> where
        C: Fn(&'b [O]) -> nom::IResult<&'b [O], CO, E> {

        match self.needed() {
            Some(nom::Needed::Size(bytes)) => try!(self.fill(bytes)),
            Some(nom::Needed::Unknown) => try!(self.fill(1)),
            None => ()
        }

        Ok(try!(self.apply(combinator)))
    }
}

#[derive(Debug)]
pub enum FillApplyError<I, E> {
    Io(io::Error),
    Parser(nom::Err<I, E>),
}

impl<I, E> From<io::Error> for FillApplyError<I, E> {
    fn from(e: io::Error) -> FillApplyError<I, E> {
        FillApplyError::Io(e)
    }
}

impl<I, E> From<nom::Err<I, E>> for FillApplyError<I, E> {
    fn from(e: nom::Err<I, E>) -> FillApplyError<I, E> {
        FillApplyError::Parser(e)
    }
}

impl<I, E> Display for FillApplyError<I, E> where I: Debug, E: Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &FillApplyError::Io(ref e) => write!(f, "Failed to read data (IO error): {}", e),
            &FillApplyError::Parser(ref e) => write!(f, "Failed to parse data (parser error): {}", e),
        }
    }
}

pub struct ReadStreamBuf<R: Read> {
    reader: R,
    //TODO: alignment?
    buf: Vec<u8>,
    cap: usize,
    needed: Cell<Option<nom::Needed>>,
    offset: Cell<usize>,
}

impl<R: Read> ReadStreamBuf<R> {
    pub fn new(reader: R) -> ReadStreamBuf<R> {
        ReadStreamBuf::with_capacity(reader, DEFAULT_BUF_SIZE)
    }

    pub fn with_capacity(reader: R, cap: usize) -> ReadStreamBuf<R> {
        ReadStreamBuf {
            reader: reader,
            buf: Vec::with_capacity(cap),
            cap: cap,
            needed: Cell::new(Some(nom::Needed::Unknown)),
            offset: Cell::new(0),
        }
    }
}

impl<R: Read> StreamBuf<u8> for ReadStreamBuf<R> {
    fn fill(&mut self, min_bytes: usize) -> Result<(), io::Error> {
        let len = self.buf.len();
        let have = len - self.offset.get();
        let needed = min_bytes - have;

        if have >= min_bytes {
            return Ok(())
        }

        //TODO: enforce cap
        self.buf.resize(len + needed, 0);
        //trace!("fill needed: {}", needed);
        //trace!("buf write: {}..{} ({}); have: {} will have: {}", len, len + needed, self.buf[len..len + needed].len(), have, have + needed);
        let result = self.reader.read_exact(&mut self.buf[len..len + needed]);

        if result.is_err() {
            self.buf.resize(len, 0);
        }

        //trace!("buf have: {:?}", self.data());

        //TODO: Try to read all that we have non-blocking in case we have more
        result
    }

    fn recycle(&mut self) {
        if self.offset.get() == 0 {
            return
        }
        self.buf = self.buf.split_off(self.offset.get());
        self.offset.set(0);
    }

    fn consume(&mut self, bytes: usize) {
        let len = self.buf.len();

        let consume = if self.offset.get() + bytes > len {
            len - self.offset.get()
        } else {
            bytes
        };
        self.offset.set(self.offset.get() + consume);
    }

    fn data<'b>(&'b self) -> &'b[u8] {
        &self.buf[self.offset.get()..self.buf.len()]
    }

    fn needed(&self) -> Option<nom::Needed> {
        self.needed.get().clone()
    }

    fn apply<'b, C, CO, E>(&'b self, combinator: C) -> Result<Option<CO>, nom::Err<&'b [u8], E>> where
        C: Fn(&'b [u8]) -> nom::IResult<&'b [u8], CO, E> {
        let data = self.data();
        match combinator(data) {
            nom::IResult::Done(left, out) => {
                let consumed = data.len() - left.len();
                //trace!("done: consumed: {}", consumed);

                // Move the offset
                self.offset.set(self.offset.get() + consumed);
                // We don't know how much we will need now
                self.needed.set(None);
                Ok(Some(out))
            },
            nom::IResult::Error(err) => Err(err),
            nom::IResult::Incomplete(needed) => {
                //trace!("incomplete: needed: {:?}", needed);
                self.needed.set(Some(needed));
                Ok(None)
            }
        }
    }
}

#[cfg(test)]
mod resd_stream_buf_tests {
    use super::StreamBuf;
    use super::ReadStreamBuf;
    use nom::{IResult, Needed};
    use std::io::Cursor;

    fn subject(data: Vec<u8>) -> ReadStreamBuf<Cursor<Vec<u8>>> {
        ReadStreamBuf::new(Cursor::new(data))
    }

    fn subject_with_default_data() -> ReadStreamBuf<Cursor<Vec<u8>>> {
        subject(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
    }

    #[test]
    fn reading() {
        let mut rsb = subject_with_default_data();
        assert_eq!(rsb.data(), [].as_ref());

        rsb.fill(2).unwrap();
        assert_eq!(rsb.data(), [0, 1].as_ref());

        rsb.fill(3).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2].as_ref());
    }

    #[test]
    fn recycle() {
        let mut rsb = subject_with_default_data();

        rsb.fill(3).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2].as_ref());

        rsb.recycle();
        assert_eq!(rsb.data(), [0, 1, 2].as_ref());

        rsb.fill(5).unwrap();
        rsb.fill(7).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4, 5, 6].as_ref());
    }

    #[test]
    fn consume() {
        let mut rsb = subject_with_default_data();

        rsb.fill(5).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4].as_ref());

        rsb.consume(2);
        assert_eq!(rsb.data(), [2, 3, 4].as_ref());
    }

    #[test]
    fn consume_more_than_we_have() {
        let mut rsb = subject_with_default_data();

        // We don't consume more than we have already in the buffor
        rsb.fill(2).unwrap();
        rsb.consume(3);
        rsb.fill(2).unwrap();
        assert_eq!(rsb.data(), [2, 3].as_ref());
    }

    #[test]
    fn apply_function() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(0));
    }

    #[test]
    fn apply_should_consume() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(0));
        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(1));
    }

    #[test]
    fn apply_converted_macro() {
        let mut rsb = subject_with_default_data();

        rsb.fill(2).unwrap();
        assert_eq!(rsb.apply(closure!(tag!([0, 1]))).unwrap(), Some([0, 1].as_ref()));
    }

    #[test]
    fn apply_closure() {
        let mut rsb = subject_with_default_data();

        rsb.fill(2).unwrap();
        assert_eq!(rsb.apply(|i| tag!(i, [0, 1])).unwrap(), Some([0, 1].as_ref()));
    }

    #[test]
    fn apply_custom_fuction_with_refs() {
        let mut rsb = subject_with_default_data();
        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }

        rsb.fill(3).unwrap();
        assert_eq!(rsb.apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
    }

    #[test]
    fn needed_with_apply() {
        let mut rsb = subject_with_default_data();

        fn comb<'a>(input: &'a[u8]) -> IResult<&'a[u8], &'a[u8]> {
            tag!(input, [0, 1, 2])
        }

        assert!(rsb.apply(comb).unwrap().is_none());

        let needed = rsb.needed().unwrap();
        assert_eq!(needed, Needed::Size(3));

        if let Needed::Size(bytes) = needed {
            rsb.fill(bytes).unwrap();
            assert_eq!(rsb.apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
        }
    }

    #[test]
    fn fill_apply() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }

        assert_eq!(rsb.fill_apply(comb).unwrap(), None);
        assert_eq!(rsb.fill_apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
        assert_eq!(rsb.fill_apply(be_u8).unwrap(), None);
        assert_eq!(rsb.fill_apply(be_u8).unwrap(), Some(3));
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::net::SocketAddr;
use std::io::{Read, Write, BufRead};

pub fn main() {
    let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string());
    let addr = addr.parse::<SocketAddr>().unwrap();

    let mut socket = ::std::net::TcpStream::connect(addr).unwrap();

    socket.write_all(&[1]).unwrap(); // subscribe

    let mut message = [0u8; 256];
    let mut header = [0u8];

    loop {
        socket.read_exact(&mut header).unwrap();
        let len = header[0] as usize;
        socket.read_exact(&mut message[..len]).unwrap();
        match ::std::str::from_utf8(&message[..len]) {
            Ok(s) => {
                println!("{}", s);
            }
            Err(e) => {
                println!("[received non-utf8 data]");
            }
        }
    }
}
<commit_msg>better eof handling<commit_after>use std::env;
use std::net::SocketAddr;
use std::io::{Read, Write};

pub fn main() {
    let addr = env::args().nth(1).unwrap_or("127.0.0.1:8080".to_string());
    let addr = addr.parse::<SocketAddr>().unwrap();

    let mut socket = ::std::net::TcpStream::connect(addr).unwrap();

    socket.write_all(&[1]).unwrap(); // subscribe

    let mut message = [0u8; 256];
    let mut header = [0u8];

    loop {
        if socket.read(&mut header).unwrap() == 0 {
            // EOF
            break;
        }

        let len = header[0] as usize;
        socket.read_exact(&mut message[..len]).unwrap();
        match ::std::str::from_utf8(&message[..len]) {
            Ok(s) => {
                println!("{}", s);
            }
            Err(_) => {
                println!("[received non-utf8 data]");
            }
        }
    }
}
<|endoftext|>


<commit_before>use {Errno, Error, Result};
use libc;
use void::Void;
use std::mem::drop;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RebootMode {
    Halt,
    kexec,
    PowerOff,
    Restart,
    // we do not support Restart2,
    Suspend,
}

pub fn reboot(how: RebootMode) -> Result<Void> {
    let cmd = match how {
        RebootMode::Halt => libc::RB_HALT_SYSTEM,
        RebootMode::kexec => libc::RB_KEXEC,
        RebootMode::PowerOff => libc::RB_POWER_OFF,
        RebootMode::Restart => libc::RB_AUTOBOOT,
        // we do not support Restart2,
        RebootMode::Suspend => libc::RB_SW_SUSPEND,
    };
    unsafe {
        libc::reboot(cmd)
    };
    Err(Error::Sys(Errno::last()))
}

/// Enable or disable the reboot keystroke (Ctrl-Alt-Delete).
///
/// Corresponds to calling `reboot(RB_ENABLE_CAD)` or `reboot(RB_DISABLE_CAD)` in C.
pub fn set_cad_enabled(enable: bool) -> Result<()> {
    let cmd = if enable {
        libc::RB_ENABLE_CAD
    } else {
        libc::RB_DISABLE_CAD
    };
    let res = unsafe {
        libc::reboot(cmd)
    };
    Errno::result(res).map(drop)
}
<commit_msg>Add some documentation<commit_after>//! Reboot/shutdown or enable/disable Ctrl-Alt-Delete.

use {Errno, Error, Result};
use libc;
use void::Void;
use std::mem::drop;

/// How exactly should the system be rebooted.
///
/// See [`set_cad_enabled()`](fn.set_cad_enabled.html) for
/// enabling/disabling Ctrl-Alt-Delete.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RebootMode {
    Halt,
    kexec,
    PowerOff,
    Restart,
    // we do not support Restart2,
    Suspend,
}

pub fn reboot(how: RebootMode) -> Result<Void> {
    let cmd = match how {
        RebootMode::Halt => libc::RB_HALT_SYSTEM,
        RebootMode::kexec => libc::RB_KEXEC,
        RebootMode::PowerOff => libc::RB_POWER_OFF,
        RebootMode::Restart => libc::RB_AUTOBOOT,
        // we do not support Restart2,
        RebootMode::Suspend => libc::RB_SW_SUSPEND,
    };
    unsafe {
        libc::reboot(cmd)
    };
    Err(Error::Sys(Errno::last()))
}

/// Enable or disable the reboot keystroke (Ctrl-Alt-Delete).
///
/// Corresponds to calling `reboot(RB_ENABLE_CAD)` or `reboot(RB_DISABLE_CAD)` in C.
pub fn set_cad_enabled(enable: bool) -> Result<()> {
    let cmd = if enable {
        libc::RB_ENABLE_CAD
    } else {
        libc::RB_DISABLE_CAD
    };
    let res = unsafe {
        libc::reboot(cmd)
    };
    Errno::result(res).map(drop)
}
<|endoftext|>