<commit_before>use std::io::{BufRead, Cursor, Read};
use std::mem::size_of;
use byteorder::{BigEndian, ReadBytesExt};

use bits::test_bit;

/// TCP Header flags.
pub struct TcpFlags {
    /// ECN-nonce concealment protection.
    pub ns: bool,
    /// Congestion window reduced.
    pub cwr: bool,
    /// ECN-Echo.
    pub ece: bool,
    /// Indicates the Urgent pointer field is significant.
    pub urg: bool,
    /// Indicates that the Acknowledgment field is significant.
    pub ack: bool,
    /// Asks to push the buffered data to the receiving application.
    pub psh: bool,
    /// Reset the connection.
    pub rst: bool,
    /// Synchronize sequence numbers.
    pub syn: bool,
    /// No more data from sender.
    pub fin: bool,
}

impl TcpFlags {
    fn of_int(d: u16) -> TcpFlags {
        TcpFlags {
            ns: test_bit(0, d as u64),
            cwr: test_bit(1, d as u64),
            ece: test_bit(2, d as u64),
            urg: test_bit(3, d as u64),
            ack: test_bit(4, d as u64),
            psh: test_bit(5, d as u64),
            rst: test_bit(6, d as u64),
            syn: test_bit(7, d as u64),
            fin: test_bit(8, d as u64),
        }
    }
}

/// TCP frame of a packet.
pub struct Tcp {
    pub src: u16,
    pub dst: u16,
    pub seq: u32,
    pub ack: u32,
    pub offset: u8,
    pub flags: TcpFlags,
    pub window: u16,
    pub chksum: u16,
    pub urgent: u16,
    pub payload: Vec<u8>,
}

#[repr(packed)]
struct TcpNet(u16, u16, u32, u32, u16, u16, u16, u16);

impl Tcp {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Tcp> {
        if bytes.get_ref().len() < size_of::<TcpNet>() {
            return None;
        }
        let src = bytes.read_u16::<BigEndian>().unwrap();
        let dst = bytes.read_u16::<BigEndian>().unwrap();
        let seq = bytes.read_u32::<BigEndian>().unwrap();
        let ack = bytes.read_u32::<BigEndian>().unwrap();
        let offset = bytes.read_u16::<BigEndian>().unwrap();
        let flags = TcpFlags::of_int(offset);
        let offset = (offset >> 12) as u8 & 0x0f;
        let window = bytes.read_u16::<BigEndian>().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let urgent = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Tcp {
            src: src,
            dst: dst,
            seq: seq,
            ack: ack,
            offset: offset,
            flags: flags,
            window: window,
            chksum: chksum,
            urgent: urgent,
            payload: payload,
        })
    }
}

/// UDP frame of a packet.
pub struct Udp {
    pub src: u16,
    pub dst: u16,
    pub chksum: u16,
    pub payload: Vec<u8>,
}

impl Udp {
    fn size_of() -> usize {
        8
    }

    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Udp> {
        if bytes.get_ref().len() < Self::size_of() {
            return None;
        }
        let src = bytes.read_u16::<BigEndian>().unwrap();
        let dst = bytes.read_u16::<BigEndian>().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Udp {
            src: src,
            dst: dst,
            chksum: chksum,
            payload: payload,
        })
    }
}

/// ICMP frame of a packet.
pub struct Icmp {
    pub typ: u8,
    pub code: u8,
    pub chksum: u16,
    pub payload: Vec<u8>,
}

impl Icmp {
    fn size_of() -> usize {
        4
    }

    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Icmp> {
        if bytes.get_ref().len() < Self::size_of() {
            return None;
        }
        let typ = bytes.read_u8().unwrap();
        let code = bytes.read_u8().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Icmp {
            typ: typ,
            code: code,
            chksum: chksum,
            payload: payload,
        })
    }
}

/// Represents packets at the transport protocol level, which are encapsulated
/// within the IPv4 payload. At present, we only support TCP, UDP, and ICMP
/// explicitly; otherwise, the raw bytes and IPv4 protocol number are provided.
pub enum Tp {
    Tcp(Tcp),
    Udp(Udp),
    Icmp(Icmp),
    Unparsable(u8, Vec<u8>),
}

/// The type of IPv4 flags.
pub struct Flags {
    pub dont_fragment: bool,
    pub more_fragments: bool,
}

impl Flags {
    fn of_int(flags: u32) -> Flags {
        Flags {
            dont_fragment: test_bit(1, flags as u64),
            more_fragments: test_bit(2, flags as u64),
        }
    }
}

/// IPv4 frame of a packet.
pub struct Ip {
    pub tos: u8,
    pub ident: u16,
    pub flags: Flags,
    pub frag: u16,
    pub ttl: u8,
    pub chksum: u16,
    pub src: u32,
    pub dst: u32,
    pub options: Vec<u8>,
    pub tp: Tp,
}

#[repr(u8)]
enum IpProto {
    IpICMP = 0x01,
    IpTCP = 0x06,
    IpUDP = 0x11,
}

#[repr(packed)]
struct IpNet(u8, u8, u16, u16, u16, u8, u8, u16, u32, u32);

impl Ip {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Ip> {
        if bytes.get_ref().len() < size_of::<IpNet>() {
            return None;
        }
        let vhl = bytes.read_u8().unwrap();
        if (vhl >> 4) != 4 {
            return None;
        }
        let ihl = vhl & 0x0f;
        let tos = bytes.read_u8().unwrap();
        bytes.consume(2);
        let ident = bytes.read_u16::<BigEndian>().unwrap();
        let frag = bytes.read_u16::<BigEndian>().unwrap();
        let flags = Flags::of_int((frag as u32) >> 13);
        let ttl = bytes.read_u8().unwrap();
        let proto = bytes.read_u8().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let src = bytes.read_u32::<BigEndian>().unwrap();
        let dst = bytes.read_u32::<BigEndian>().unwrap();
        let options_len = (ihl * 4) as usize - size_of::<IpNet>();
        let mut options = vec![0; options_len];
        bytes.read_exact(&mut options).unwrap();
        let tp = match proto {
            t if t == (IpProto::IpICMP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let icmp = Icmp::parse(bytes);
                if icmp.is_some() {
                    Tp::Icmp(icmp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            t if t == (IpProto::IpTCP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let tcp = Tcp::parse(bytes);
                if tcp.is_some() {
                    Tp::Tcp(tcp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            t if t == (IpProto::IpUDP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let udp = Udp::parse(bytes);
                if udp.is_some() {
                    Tp::Udp(udp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            _ => Tp::Unparsable(proto, bytes.get_ref().clone()),
        };
        Some(Ip {
            tos: tos,
            ident: ident,
            flags: flags,
            frag: frag,
            ttl: ttl,
            chksum: chksum,
            src: src,
            dst: dst,
            options: options,
            tp: tp,
        })
    }
}

/// Address resolution protocol (ARP) packet payload.
pub enum Arp {
    Query([u8; 6], u32, u32),
    Reply([u8; 6], u32, [u8; 6], u32),
}

#[repr(packed)]
struct ArpNet(u16, u16, u8, u8, u16, [u8; 6], u32, [u8; 6], u32);

impl Arp {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Arp> {
        if bytes.get_ref().len() < size_of::<ArpNet>() {
            return None;
        }
        bytes.consume(6);
        let oper = bytes.read_u16::<BigEndian>().unwrap();
        let mut sha: [u8; 6] = [0; 6];
        for i in 0..6 {
            sha[i] = bytes.read_u8().unwrap();
        }
        let spa = bytes.read_u32::<BigEndian>().unwrap();
        let mut tha: [u8; 6] = [0; 6];
        for i in 0..6 {
            tha[i] = bytes.read_u8().unwrap();
        }
        let tpa = bytes.read_u32::<BigEndian>().unwrap();
        match oper {
            0x0001 => Some(Arp::Query(sha, spa, tpa)),
            0x0002 => Some(Arp::Reply(sha, spa, tha, tpa)),
            _ => None,
        }
    }
}

/// Represents a packet at the network protocol level.
pub enum Nw {
    Ip(Ip),
    Arp(Arp),
    Unparsable(u16, Vec<u8>),
}

/// Represents a packet at the ethernet protocol level.
pub struct Packet {
    pub dl_src: [u8; 6],
    pub dl_dst: [u8; 6],
    pub dl_vlan: Option<u16>,
    pub dl_vlan_dei: bool,
    pub dl_vlan_pcp: u8,
    pub nw: Nw,
}

#[repr(u16)]
enum EthTyp {
    EthTypIP = 0x0800,
    EthTypARP = 0x0806,
    EthTypVLAN = 0x8100,
}

impl Packet {
    pub fn parse(buf: &[u8]) -> Packet {
        let mut bytes = Cursor::new(buf.to_vec());
        let mut dst: [u8; 6] = [0; 6];
        let mut src: [u8; 6] = [0; 6];
        for i in 0..6 {
            dst[i] = bytes.read_u8().unwrap();
        }
        for i in 0..6 {
            src[i] = bytes.read_u8().unwrap();
        }
        let typ = bytes.read_u16::<BigEndian>().unwrap();
        let (tag, dei, pcp, typ) = match typ {
            t if t == (EthTyp::EthTypVLAN as u16) => {
                let tag_and_pcp = bytes.read_u16::<BigEndian>().unwrap();
                let tag = tag_and_pcp & 0xfff;
                let dei = (tag_and_pcp & 0x1000) > 0;
                let pcp = tag_and_pcp >> 13;
                let typ = bytes.read_u16::<BigEndian>().unwrap();
                (Some(tag), dei, pcp as u8, typ)
            }
            _ => (None, false, 0x0, typ),
        };
        let nw_header = match typ {
            t if t == (EthTyp::EthTypIP as u16) => {
                let bytes_ = bytes.get_ref().clone();
                let ip = Ip::parse(&mut bytes);
                if ip.is_some() {
                    Nw::Ip(ip.unwrap())
                } else {
                    Nw::Unparsable(typ, bytes_)
                }
            }
            t if t == (EthTyp::EthTypARP as u16) => {
                let bytes_ = bytes.get_ref().clone();
                let arp = Arp::parse(&mut bytes);
                if arp.is_some() {
                    Nw::Arp(arp.unwrap())
                } else {
                    Nw::Unparsable(typ, bytes_)
                }
            }
            _ => Nw::Unparsable(typ, bytes.into_inner()),
        };
        Packet {
            dl_src: src,
            dl_dst: dst,
            dl_vlan: tag,
            dl_vlan_dei: dei,
            dl_vlan_pcp: pcp,
            nw: nw_header,
        }
    }
}
<commit_msg>Implement mac translations<commit_after>use std::io::{BufRead, Cursor, Read};
use std::mem::size_of;
use byteorder::{BigEndian, ReadBytesExt};

use bits::test_bit;

pub fn bytes_of_mac(addr: u64) -> [u8; 6] {
    let mut arr = [0; 6];
    for i in 0..6 {
        let iso = 0xff0000000000 >> i;

        arr[i] = ((iso & addr) >> (5 - i)) as u8;
    }
    arr
}

pub fn mac_of_bytes(addr: [u8; 6]) -> u64 {
    fn byte(u: &[u8; 6], i: usize) -> u64 {
        u[i] as u64
    };
    (byte(&addr, 0) << 8 * 5) | (byte(&addr, 1) << 8 * 4) | (byte(&addr, 2) << 8 * 3) |
    (byte(&addr, 3) << 8 * 2) | (byte(&addr, 4) << 8 * 1) | (byte(&addr, 5))
}

/// TCP Header flags.
pub struct TcpFlags {
    /// ECN-nonce concealment protection.
    pub ns: bool,
    /// Congestion window reduced.
    pub cwr: bool,
    /// ECN-Echo.
    pub ece: bool,
    /// Indicates the Urgent pointer field is significant.
    pub urg: bool,
    /// Indicates that the Acknowledgment field is significant.
    pub ack: bool,
    /// Asks to push the buffered data to the receiving application.
    pub psh: bool,
    /// Reset the connection.
    pub rst: bool,
    /// Synchronize sequence numbers.
    pub syn: bool,
    /// No more data from sender.
    pub fin: bool,
}

impl TcpFlags {
    fn of_int(d: u16) -> TcpFlags {
        TcpFlags {
            ns: test_bit(0, d as u64),
            cwr: test_bit(1, d as u64),
            ece: test_bit(2, d as u64),
            urg: test_bit(3, d as u64),
            ack: test_bit(4, d as u64),
            psh: test_bit(5, d as u64),
            rst: test_bit(6, d as u64),
            syn: test_bit(7, d as u64),
            fin: test_bit(8, d as u64),
        }
    }
}

/// TCP frame of a packet.
pub struct Tcp {
    pub src: u16,
    pub dst: u16,
    pub seq: u32,
    pub ack: u32,
    pub offset: u8,
    pub flags: TcpFlags,
    pub window: u16,
    pub chksum: u16,
    pub urgent: u16,
    pub payload: Vec<u8>,
}

#[repr(packed)]
struct TcpNet(u16, u16, u32, u32, u16, u16, u16, u16);

impl Tcp {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Tcp> {
        if bytes.get_ref().len() < size_of::<TcpNet>() {
            return None;
        }
        let src = bytes.read_u16::<BigEndian>().unwrap();
        let dst = bytes.read_u16::<BigEndian>().unwrap();
        let seq = bytes.read_u32::<BigEndian>().unwrap();
        let ack = bytes.read_u32::<BigEndian>().unwrap();
        let offset = bytes.read_u16::<BigEndian>().unwrap();
        let flags = TcpFlags::of_int(offset);
        let offset = (offset >> 12) as u8 & 0x0f;
        let window = bytes.read_u16::<BigEndian>().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let urgent = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Tcp {
            src: src,
            dst: dst,
            seq: seq,
            ack: ack,
            offset: offset,
            flags: flags,
            window: window,
            chksum: chksum,
            urgent: urgent,
            payload: payload,
        })
    }
}

/// UDP frame of a packet.
pub struct Udp {
    pub src: u16,
    pub dst: u16,
    pub chksum: u16,
    pub payload: Vec<u8>,
}

impl Udp {
    fn size_of() -> usize {
        8
    }

    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Udp> {
        if bytes.get_ref().len() < Self::size_of() {
            return None;
        }
        let src = bytes.read_u16::<BigEndian>().unwrap();
        let dst = bytes.read_u16::<BigEndian>().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Udp {
            src: src,
            dst: dst,
            chksum: chksum,
            payload: payload,
        })
    }
}

/// ICMP frame of a packet.
pub struct Icmp {
    pub typ: u8,
    pub code: u8,
    pub chksum: u16,
    pub payload: Vec<u8>,
}

impl Icmp {
    fn size_of() -> usize {
        4
    }

    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Icmp> {
        if bytes.get_ref().len() < Self::size_of() {
            return None;
        }
        let typ = bytes.read_u8().unwrap();
        let code = bytes.read_u8().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let mut payload = vec![0; bytes.get_ref().len()];
        bytes.read_exact(&mut payload).unwrap();
        Some(Icmp {
            typ: typ,
            code: code,
            chksum: chksum,
            payload: payload,
        })
    }
}

/// Represents packets at the transport protocol level, which are encapsulated
/// within the IPv4 payload. At present, we only support TCP, UDP, and ICMP
/// explicitly; otherwise, the raw bytes and IPv4 protocol number are provided.
pub enum Tp {
    Tcp(Tcp),
    Udp(Udp),
    Icmp(Icmp),
    Unparsable(u8, Vec<u8>),
}

/// The type of IPv4 flags.
pub struct Flags {
    pub dont_fragment: bool,
    pub more_fragments: bool,
}

impl Flags {
    fn of_int(flags: u32) -> Flags {
        Flags {
            dont_fragment: test_bit(1, flags as u64),
            more_fragments: test_bit(2, flags as u64),
        }
    }
}

/// IPv4 frame of a packet.
pub struct Ip {
    pub tos: u8,
    pub ident: u16,
    pub flags: Flags,
    pub frag: u16,
    pub ttl: u8,
    pub chksum: u16,
    pub src: u32,
    pub dst: u32,
    pub options: Vec<u8>,
    pub tp: Tp,
}

#[repr(u8)]
enum IpProto {
    IpICMP = 0x01,
    IpTCP = 0x06,
    IpUDP = 0x11,
}

#[repr(packed)]
struct IpNet(u8, u8, u16, u16, u16, u8, u8, u16, u32, u32);

impl Ip {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Ip> {
        if bytes.get_ref().len() < size_of::<IpNet>() {
            return None;
        }
        let vhl = bytes.read_u8().unwrap();
        if (vhl >> 4) != 4 {
            return None;
        }
        let ihl = vhl & 0x0f;
        let tos = bytes.read_u8().unwrap();
        bytes.consume(2);
        let ident = bytes.read_u16::<BigEndian>().unwrap();
        let frag = bytes.read_u16::<BigEndian>().unwrap();
        let flags = Flags::of_int((frag as u32) >> 13);
        let ttl = bytes.read_u8().unwrap();
        let proto = bytes.read_u8().unwrap();
        let chksum = bytes.read_u16::<BigEndian>().unwrap();
        let src = bytes.read_u32::<BigEndian>().unwrap();
        let dst = bytes.read_u32::<BigEndian>().unwrap();
        let options_len = (ihl * 4) as usize - size_of::<IpNet>();
        let mut options = vec![0; options_len];
        bytes.read_exact(&mut options).unwrap();
        let tp = match proto {
            t if t == (IpProto::IpICMP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let icmp = Icmp::parse(bytes);
                if icmp.is_some() {
                    Tp::Icmp(icmp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            t if t == (IpProto::IpTCP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let tcp = Tcp::parse(bytes);
                if tcp.is_some() {
                    Tp::Tcp(tcp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            t if t == (IpProto::IpUDP as u8) => {
                let bytes_ = bytes.get_ref().clone();
                let udp = Udp::parse(bytes);
                if udp.is_some() {
                    Tp::Udp(udp.unwrap())
                } else {
                    Tp::Unparsable(proto, bytes_)
                }
            }
            _ => Tp::Unparsable(proto, bytes.get_ref().clone()),
        };
        Some(Ip {
            tos: tos,
            ident: ident,
            flags: flags,
            frag: frag,
            ttl: ttl,
            chksum: chksum,
            src: src,
            dst: dst,
            options: options,
            tp: tp,
        })
    }
}

/// Address resolution protocol (ARP) packet payload.
pub enum Arp {
    Query([u8; 6], u32, u32),
    Reply([u8; 6], u32, [u8; 6], u32),
}

#[repr(packed)]
struct ArpNet(u16, u16, u8, u8, u16, [u8; 6], u32, [u8; 6], u32);

impl Arp {
    fn parse(bytes: &mut Cursor<Vec<u8>>) -> Option<Arp> {
        if bytes.get_ref().len() < size_of::<ArpNet>() {
            return None;
        }
        bytes.consume(6);
        let oper = bytes.read_u16::<BigEndian>().unwrap();
        let mut sha: [u8; 6] = [0; 6];
        for i in 0..6 {
            sha[i] = bytes.read_u8().unwrap();
        }
        let spa = bytes.read_u32::<BigEndian>().unwrap();
        let mut tha: [u8; 6] = [0; 6];
        for i in 0..6 {
            tha[i] = bytes.read_u8().unwrap();
        }
        let tpa = bytes.read_u32::<BigEndian>().unwrap();
        match oper {
            0x0001 => Some(Arp::Query(sha, spa, tpa)),
            0x0002 => Some(Arp::Reply(sha, spa, tha, tpa)),
            _ => None,
        }
    }
}

/// Represents a packet at the network protocol level.
pub enum Nw {
    Ip(Ip),
    Arp(Arp),
    Unparsable(u16, Vec<u8>),
}

/// Represents a packet at the ethernet protocol level.
pub struct Packet {
    pub dl_src: [u8; 6],
    pub dl_dst: [u8; 6],
    pub dl_vlan: Option<u16>,
    pub dl_vlan_dei: bool,
    pub dl_vlan_pcp: u8,
    pub nw: Nw,
}

#[repr(u16)]
enum EthTyp {
    EthTypIP = 0x0800,
    EthTypARP = 0x0806,
    EthTypVLAN = 0x8100,
}

impl Packet {
    pub fn parse(buf: &[u8]) -> Packet {
        let mut bytes = Cursor::new(buf.to_vec());
        let mut dst: [u8; 6] = [0; 6];
        let mut src: [u8; 6] = [0; 6];
        for i in 0..6 {
            dst[i] = bytes.read_u8().unwrap();
        }
        for i in 0..6 {
            src[i] = bytes.read_u8().unwrap();
        }
        let typ = bytes.read_u16::<BigEndian>().unwrap();
        let (tag, dei, pcp, typ) = match typ {
            t if t == (EthTyp::EthTypVLAN as u16) => {
                let tag_and_pcp = bytes.read_u16::<BigEndian>().unwrap();
                let tag = tag_and_pcp & 0xfff;
                let dei = (tag_and_pcp & 0x1000) > 0;
                let pcp = tag_and_pcp >> 13;
                let typ = bytes.read_u16::<BigEndian>().unwrap();
                (Some(tag), dei, pcp as u8, typ)
            }
            _ => (None, false, 0x0, typ),
        };
        let nw_header = match typ {
            t if t == (EthTyp::EthTypIP as u16) => {
                let bytes_ = bytes.get_ref().clone();
                let ip = Ip::parse(&mut bytes);
                if ip.is_some() {
                    Nw::Ip(ip.unwrap())
                } else {
                    Nw::Unparsable(typ, bytes_)
                }
            }
            t if t == (EthTyp::EthTypARP as u16) => {
                let bytes_ = bytes.get_ref().clone();
                let arp = Arp::parse(&mut bytes);
                if arp.is_some() {
                    Nw::Arp(arp.unwrap())
                } else {
                    Nw::Unparsable(typ, bytes_)
                }
            }
            _ => Nw::Unparsable(typ, bytes.into_inner()),
        };
        Packet {
            dl_src: src,
            dl_dst: dst,
            dl_vlan: tag,
            dl_vlan_dei: dei,
            dl_vlan_pcp: pcp,
            nw: nw_header,
        }
    }
}
<|endoftext|>


<commit_before>pub use self::grammar::sequence_set;

peg_file! grammar("grammar.rustpeg")

#[deriving(Clone, PartialEq, Show)]
enum SequenceItem {
    Number(uint),
    Range(Box<SequenceItem>, Box<SequenceItem>),
    All
}

#[cfg(test)]
mod tests {
    use super::sequence_set;
    use super::{All, Number, Range};

    #[test]
    fn test_invalid_sequences() {
        assert!(sequence_set("").is_err());
        assert!(sequence_set("a").is_err());
        assert!(sequence_set("0").is_err());
        assert!(sequence_set("a:*").is_err());
        assert!(sequence_set(":*").is_err());
        assert!(sequence_set("1:").is_err());
        assert!(sequence_set("1:0").is_err());
        assert!(sequence_set("0:1").is_err());
        assert!(sequence_set("4,5,6,").is_err());
    }

    #[test]
    fn test_sequence_num() {
        let seq = sequence_set("4324");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![Number(4324)];
        assert_eq!(seq, expected);
    }

    #[test]
    fn test_sequence_all() {
        let seq = sequence_set("*");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![All];
        assert_eq!(seq, expected);
    }

    #[test]
    fn test_sequence_ranges() {
        let seq = sequence_set("98:100");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![Range(box Number(98), box Number(100))];
        assert_eq!(seq, expected);

        assert!(sequence_set("1:5").is_ok());
        assert!(sequence_set("21:44").is_ok());
    }

    #[test]
    fn test_sequence_range_all() {
        let seq = sequence_set("31:*");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![Range(box Number(31), box All)];
        assert_eq!(seq, expected);
    }

    #[test]
    fn test_sequence_set() {
        let seq = sequence_set("1231,1342,12,98:104,16");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![Number(1231), Number(1342), Number(12), Range(box Number(98), box Number(104)), Number(16)];
        assert_eq!(seq, expected);
    }
}

<commit_msg>Remove trailing whitespace<commit_after>pub use self::grammar::sequence_set;

peg_file! grammar("grammar.rustpeg")

#[deriving(Clone, PartialEq, Show)]
enum SequenceItem {
    Number(uint),
    Range(Box<SequenceItem>, Box<SequenceItem>),
    All
}

#[cfg(test)]
mod tests {
    use super::sequence_set;
    use super::{All, Number, Range};

    #[test]
    fn test_invalid_sequences() {
        assert!(sequence_set("").is_err());
        assert!(sequence_set("a").is_err());
        assert!(sequence_set("0").is_err());
        assert!(sequence_set("a:*").is_err());
        assert!(sequence_set(":*").is_err());
        assert!(sequence_set("1:").is_err());
        assert!(sequence_set("1:0").is_err());
        assert!(sequence_set("0:1").is_err());
        assert!(sequence_set("4,5,6,").is_err());
    }

    #[test]
    fn test_sequence_num() {
        let seq = sequence_set("4324");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![Number(4324)];
        assert_eq!(seq, expected);
    }

    #[test]
    fn test_sequence_all() {
        let seq = sequence_set("*");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![All];
        assert_eq!(seq, expected);
    }

    #[test]
    fn test_sequence_ranges() {
        let seq = sequence_set("98:100");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![Range(box Number(98), box Number(100))];
        assert_eq!(seq, expected);

        assert!(sequence_set("1:5").is_ok());
        assert!(sequence_set("21:44").is_ok());
    }

    #[test]
    fn test_sequence_range_all() {
        let seq = sequence_set("31:*");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![Range(box Number(31), box All)];
        assert_eq!(seq, expected);
    }

    #[test]
    fn test_sequence_set() {
        let seq = sequence_set("1231,1342,12,98:104,16");
        assert!(seq.is_ok());
        let seq = seq.unwrap();
        let expected = vec![Number(1231), Number(1342), Number(12), Range(box Number(98), box Number(104)), Number(16)];
        assert_eq!(seq, expected);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2018 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

#![cfg_attr(feature = "cargo-clippy", allow(duration_subsec))]

use std::io;
use std::result;
use std::time::Duration;

use libc;

pub use libc::{epoll_event, EPOLLERR, EPOLLET, EPOLLIN, EPOLLONESHOT, EPOLLOUT, EPOLLPRI};

pub type Result<T> = result::Result<T, io::Error>;

fn parse_retval(retval: libc::c_int) -> Result<i32> {
    if retval == -1 {
        Err(io::Error::last_os_error())
    } else {
        Ok(retval)
    }
}

// We're using EventFd to wake up another thread
// that's waiting for epoll_wait() to return.
#[derive(Debug)]
pub struct EventFd {
    fd: i32,
}

impl EventFd {
    pub fn new() -> Result<EventFd> {
        Ok(EventFd {
            fd: unsafe {
                parse_retval(libc::eventfd(0, libc::EFD_NONBLOCK | libc::EFD_SEMAPHORE))?
            },
        })
    }

    pub fn notify(&self) -> Result<()> {
        let buffer: u64 = 1;

        if unsafe { libc::write(self.fd, &buffer as *const u64 as *const libc::c_void, 8) } == -1 {
            Err(io::Error::last_os_error())
        } else {
            Ok(())
        }
    }

    pub fn fd(&self) -> i32 {
        self.fd
    }
}

impl Drop for EventFd {
    fn drop(&mut self) {
        unsafe {
            libc::close(self.fd);
        }
    }
}

#[derive(Debug)]
pub struct Epoll {
    fd: libc::c_int,
}

impl Epoll {
    pub fn new() -> Result<Epoll> {
        Ok(Epoll {
            fd: unsafe { parse_retval(libc::epoll_create1(0))? },
        })
    }

    pub fn add(&self, fd: i32, id: u64, event_mask: i32) -> Result<()> {
        let mut event = libc::epoll_event {
            events: event_mask as u32,
            u64: id as u64,
        };

        unsafe {
            parse_retval(libc::epoll_ctl(
                self.fd,
                libc::EPOLL_CTL_ADD,
                fd,
                &mut event,
            ))?;
        }

        Ok(())
    }

    pub fn modify(&self, fd: i32, id: u64, event_mask: i32) -> Result<()> {
        let mut event = libc::epoll_event {
            events: event_mask as u32,
            u64: id as u64,
        };

        unsafe {
            parse_retval(libc::epoll_ctl(
                self.fd,
                libc::EPOLL_CTL_MOD,
                fd,
                &mut event,
            ))?;
        }

        Ok(())
    }

    pub fn delete(&self, fd: i32) -> Result<()> {
        let mut event = libc::epoll_event { events: 0, u64: 0 };

        unsafe {
            parse_retval(libc::epoll_ctl(
                self.fd,
                libc::EPOLL_CTL_DEL,
                fd,
                &mut event,
            ))?;
        }

        Ok(())
    }

    pub fn wait(
        &self,
        events: &mut [libc::epoll_event],
        timeout: Option<Duration>,
    ) -> Result<usize> {
        if events.is_empty() {
            return Ok(0);
        }

        let timeout: i32 = if let Some(duration) = timeout {
            (duration.as_secs() * 1_000) as i32 + (duration.subsec_nanos() / 1_000_000) as i32
            // (duration.as_secs() * 1_000) as i32 + duration.subsec_millis() as i32
        } else {
            -1
        };

        unsafe {
            Ok(parse_retval(libc::epoll_wait(
                self.fd,
                events.as_mut_ptr(),
                events.len() as i32,
                timeout,
            ))? as usize)
        }
    }
}

impl Drop for Epoll {
    fn drop(&mut self) {
        unsafe {
            libc::close(self.fd);
        }
    }
}
<commit_msg>Switch back to subsec_millis()<commit_after>// Copyright (c) 2017-2018 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

#![cfg_attr(feature = "cargo-clippy", allow(duration_subsec))]

use std::io;
use std::result;
use std::time::Duration;

use libc;

pub use libc::{epoll_event, EPOLLERR, EPOLLET, EPOLLIN, EPOLLONESHOT, EPOLLOUT, EPOLLPRI};

pub type Result<T> = result::Result<T, io::Error>;

fn parse_retval(retval: libc::c_int) -> Result<i32> {
    if retval == -1 {
        Err(io::Error::last_os_error())
    } else {
        Ok(retval)
    }
}

// We're using EventFd to wake up another thread
// that's waiting for epoll_wait() to return.
#[derive(Debug)]
pub struct EventFd {
    fd: i32,
}

impl EventFd {
    pub fn new() -> Result<EventFd> {
        Ok(EventFd {
            fd: unsafe {
                parse_retval(libc::eventfd(0, libc::EFD_NONBLOCK | libc::EFD_SEMAPHORE))?
            },
        })
    }

    pub fn notify(&self) -> Result<()> {
        let buffer: u64 = 1;

        if unsafe { libc::write(self.fd, &buffer as *const u64 as *const libc::c_void, 8) } == -1 {
            Err(io::Error::last_os_error())
        } else {
            Ok(())
        }
    }

    pub fn fd(&self) -> i32 {
        self.fd
    }
}

impl Drop for EventFd {
    fn drop(&mut self) {
        unsafe {
            libc::close(self.fd);
        }
    }
}

#[derive(Debug)]
pub struct Epoll {
    fd: libc::c_int,
}

impl Epoll {
    pub fn new() -> Result<Epoll> {
        Ok(Epoll {
            fd: unsafe { parse_retval(libc::epoll_create1(0))? },
        })
    }

    pub fn add(&self, fd: i32, id: u64, event_mask: i32) -> Result<()> {
        let mut event = libc::epoll_event {
            events: event_mask as u32,
            u64: id as u64,
        };

        unsafe {
            parse_retval(libc::epoll_ctl(
                self.fd,
                libc::EPOLL_CTL_ADD,
                fd,
                &mut event,
            ))?;
        }

        Ok(())
    }

    pub fn modify(&self, fd: i32, id: u64, event_mask: i32) -> Result<()> {
        let mut event = libc::epoll_event {
            events: event_mask as u32,
            u64: id as u64,
        };

        unsafe {
            parse_retval(libc::epoll_ctl(
                self.fd,
                libc::EPOLL_CTL_MOD,
                fd,
                &mut event,
            ))?;
        }

        Ok(())
    }

    pub fn delete(&self, fd: i32) -> Result<()> {
        let mut event = libc::epoll_event { events: 0, u64: 0 };

        unsafe {
            parse_retval(libc::epoll_ctl(
                self.fd,
                libc::EPOLL_CTL_DEL,
                fd,
                &mut event,
            ))?;
        }

        Ok(())
    }

    pub fn wait(
        &self,
        events: &mut [libc::epoll_event],
        timeout: Option<Duration>,
    ) -> Result<usize> {
        if events.is_empty() {
            return Ok(0);
        }

        let timeout: i32 = if let Some(duration) = timeout {
            (duration.as_secs() * 1_000) as i32 + duration.subsec_millis() as i32
        } else {
            -1
        };

        unsafe {
            Ok(parse_retval(libc::epoll_wait(
                self.fd,
                events.as_mut_ptr(),
                events.len() as i32,
                timeout,
            ))? as usize)
        }
    }
}

impl Drop for Epoll {
    fn drop(&mut self) {
        unsafe {
            libc::close(self.fd);
        }
    }
}
<|endoftext|>


<commit_before>use std::sync::{Arc, Mutex};
use std::collections::{HashMap, HashSet, VecDeque};

use bytes::Bytes;

use job::Job;
use worker::Worker;

pub type JobQueue = VecDeque<Arc<Job>>;
pub type JobQueues = HashMap<Bytes, [JobQueue; 3]>;

pub struct JobStorage {
    jobs: HashMap<Bytes, Arc<Job>>, // Owns the job objects forever
    queues: JobQueues,
    remotes_by_unique: HashMap<Bytes, HashSet<usize>>,
    remotes_by_handle: HashMap<Bytes, Vec<usize>>,
}

pub type SharedJobStorage = Arc<Mutex<JobStorage>>;

pub trait HandleJobStorage {
    fn new_job_storage() -> SharedJobStorage;
    fn coalesce_unique(&mut self, unique: &Bytes, remote: Option<usize>) -> Option<Bytes>;
    fn add_job(&mut self, job: Arc<Job>, priority: JobQueuePriority, remote: Option<usize>);
    fn get_job(&mut self, worker: &mut Worker) -> Option<Arc<Job>>;
}

pub type JobQueuePriority = usize;

const INIT_JOB_STORAGE_CAPACITY: usize = 10000000; // XXX This should be configurable
const INIT_JOB_FUNCTIONS_CAPACITY: usize = 4096; // XXX This should be configurable
const INIT_JOB_REMOTES_CAPACITY: usize = 8;

impl JobStorage {
    fn new() -> JobStorage {
        JobStorage {
            jobs: HashMap::with_capacity(INIT_JOB_STORAGE_CAPACITY),
            queues: HashMap::with_capacity(INIT_JOB_FUNCTIONS_CAPACITY),
            remotes_by_unique: HashMap::with_capacity(INIT_JOB_STORAGE_CAPACITY),
            remotes_by_handle: HashMap::with_capacity(INIT_JOB_STORAGE_CAPACITY),
        }
    }

    pub fn queues(&self) -> &JobQueues {
        &self.queues
    }

    pub fn remove_job(&mut self, unique: &Bytes) {
        match self.jobs.get(unique) {
            None => {}
            Some(job) => {
                self.remotes_by_handle.remove(&job.handle);
            }
        }
        self.jobs.remove(unique);
        self.remotes_by_unique.remove(unique);
    }

    pub fn remotes_by_unique(&self, unique: &Bytes) -> Option<&HashSet<usize>> {
        self.remotes_by_unique.get(unique)
    }

    pub fn remotes_by_handle(&self, handle: &Bytes) -> Option<&Vec<usize>> {
        self.remotes_by_handle.get(handle)
    }
}

impl HandleJobStorage for SharedJobStorage {
    fn new_job_storage() -> SharedJobStorage {
        Arc::new(Mutex::new(JobStorage::new()))
    }

    fn coalesce_unique(&mut self, unique: &Bytes, remote: Option<usize>) -> Option<Bytes> {
        let mut storage = self.lock().unwrap();
        let handle = match storage.jobs.get(unique) {
            None => return None,
            Some(job) => job.handle.clone(),
        };
        let mut add_remote = false;
        match storage.remotes_by_unique.get_mut(unique) {
            None => warn!("Job with no remote storage found: {:?}", &handle),
            Some(remotes) => {
                match remote {
                    None => {}
                    Some(remote) => {
                        add_remote = remotes.insert(remote);
                    }
                };
            }
        }
        // We don't need two hashsets for the same set, so just push if we added to the unique set
        if add_remote {
            match storage.remotes_by_handle.get_mut(&handle) {
                None => warn!("Job with no remote storage found: {:?}", &handle),
                Some(remotes) => {
                    match remote {
                        None => {}
                        Some(remote) => {
                            remotes.push(remote);
                        }
                    };
                }
            }
        }
        Some(handle)
    }

    fn add_job(&mut self, job: Arc<Job>, priority: JobQueuePriority, remote: Option<usize>) {
        let mut storage = self.lock().unwrap();
        if !storage.queues.contains_key(&job.fname) {
            let high_queue = VecDeque::new();
            let norm_queue = VecDeque::new();
            let low_queue = VecDeque::new();
            storage.queues.insert(job.fname.clone(), [high_queue, norm_queue, low_queue]);
        }
        storage.jobs.insert(job.unique.clone(), job.clone());
        let mut remotes_by_unique = HashSet::with_capacity(INIT_JOB_REMOTES_CAPACITY);
        let mut remotes_by_handle = Vec::with_capacity(INIT_JOB_REMOTES_CAPACITY);
        match remote {
            None => {}
            Some(remote) => {
                remotes_by_unique.insert(remote);
                remotes_by_handle.push(remote);
            }
        }
        storage.remotes_by_unique.insert(job.unique.clone(), remotes_by_unique);
        storage.remotes_by_handle.insert(job.handle.clone(), remotes_by_handle);
        storage.queues.get_mut(&job.fname).unwrap()[priority].push_back(job.clone());
    }

    fn get_job(&mut self, worker: &mut Worker) -> Option<Arc<Job>> {
        let mut storage = self.lock().unwrap();
        let mut job: Option<Arc<Job>> = None;
        debug!("{:?}", &worker);
        for func in worker.iter() {
            debug!("func = {:?}", &func);
            match storage.queues.get_mut(&func) {
                None => {}
                Some(prios) => {
                    debug!("found func with {} priority queues", prios.len());
                    let mut i = 0;
                    for q in prios {
                        debug!("searching priority {}", i);
                        i = i + 1;
                        if !q.is_empty() {
                            debug!("Queue has items! {:?}", q.len());
                            job = q.pop_front();
                            break;
                        }
                    }
                }
            }
        }
        match job {
            Some(job) => {
                worker.assign_job(&job);
                Some(job)
            }
            None => None,
        }
    }
}
<commit_msg>Do one less lookup in add_job<commit_after>use std::sync::{Arc, Mutex};
use std::collections::{HashMap, HashSet, VecDeque};

use bytes::Bytes;

use job::Job;
use worker::Worker;

pub type JobQueue = VecDeque<Arc<Job>>;
pub type JobQueues = HashMap<Bytes, [JobQueue; 3]>;

pub struct JobStorage {
    jobs: HashMap<Bytes, Arc<Job>>, // Owns the job objects forever
    queues: JobQueues,
    remotes_by_unique: HashMap<Bytes, HashSet<usize>>,
    remotes_by_handle: HashMap<Bytes, Vec<usize>>,
}

pub type SharedJobStorage = Arc<Mutex<JobStorage>>;

pub trait HandleJobStorage {
    fn new_job_storage() -> SharedJobStorage;
    fn coalesce_unique(&mut self, unique: &Bytes, remote: Option<usize>) -> Option<Bytes>;
    fn add_job(&mut self, job: Arc<Job>, priority: JobQueuePriority, remote: Option<usize>);
    fn get_job(&mut self, worker: &mut Worker) -> Option<Arc<Job>>;
}

pub type JobQueuePriority = usize;

const INIT_JOB_STORAGE_CAPACITY: usize = 10000000; // XXX This should be configurable
const INIT_JOB_FUNCTIONS_CAPACITY: usize = 4096; // XXX This should be configurable
const INIT_JOB_REMOTES_CAPACITY: usize = 8;

impl JobStorage {
    fn new() -> JobStorage {
        JobStorage {
            jobs: HashMap::with_capacity(INIT_JOB_STORAGE_CAPACITY),
            queues: HashMap::with_capacity(INIT_JOB_FUNCTIONS_CAPACITY),
            remotes_by_unique: HashMap::with_capacity(INIT_JOB_STORAGE_CAPACITY),
            remotes_by_handle: HashMap::with_capacity(INIT_JOB_STORAGE_CAPACITY),
        }
    }

    pub fn queues(&self) -> &JobQueues {
        &self.queues
    }

    pub fn remove_job(&mut self, unique: &Bytes) {
        match self.jobs.get(unique) {
            None => {}
            Some(job) => {
                self.remotes_by_handle.remove(&job.handle);
            }
        }
        self.jobs.remove(unique);
        self.remotes_by_unique.remove(unique);
    }

    pub fn remotes_by_unique(&self, unique: &Bytes) -> Option<&HashSet<usize>> {
        self.remotes_by_unique.get(unique)
    }

    pub fn remotes_by_handle(&self, handle: &Bytes) -> Option<&Vec<usize>> {
        self.remotes_by_handle.get(handle)
    }
}

impl HandleJobStorage for SharedJobStorage {
    fn new_job_storage() -> SharedJobStorage {
        Arc::new(Mutex::new(JobStorage::new()))
    }

    fn coalesce_unique(&mut self, unique: &Bytes, remote: Option<usize>) -> Option<Bytes> {
        let mut storage = self.lock().unwrap();
        let handle = match storage.jobs.get(unique) {
            None => return None,
            Some(job) => job.handle.clone(),
        };
        let mut add_remote = false;
        match storage.remotes_by_unique.get_mut(unique) {
            None => warn!("Job with no remote storage found: {:?}", &handle),
            Some(remotes) => {
                match remote {
                    None => {}
                    Some(remote) => {
                        add_remote = remotes.insert(remote);
                    }
                };
            }
        }
        // We don't need two hashsets for the same set, so just push if we added to the unique set
        if add_remote {
            match storage.remotes_by_handle.get_mut(&handle) {
                None => warn!("Job with no remote storage found: {:?}", &handle),
                Some(remotes) => {
                    match remote {
                        None => {}
                        Some(remote) => {
                            remotes.push(remote);
                        }
                    };
                }
            }
        }
        Some(handle)
    }

    fn add_job(&mut self, job: Arc<Job>, priority: JobQueuePriority, remote: Option<usize>) {
        let mut storage = self.lock().unwrap();
        {
            let func_queues = storage.queues.entry(job.fname.clone()).or_insert_with(|| {
                let high_queue = VecDeque::new();
                let norm_queue = VecDeque::new();
                let low_queue = VecDeque::new();
                [high_queue, norm_queue, low_queue]
            });
            func_queues[priority].push_back(job.clone());
        }
        let mut remotes_by_unique = HashSet::with_capacity(INIT_JOB_REMOTES_CAPACITY);
        let mut remotes_by_handle = Vec::with_capacity(INIT_JOB_REMOTES_CAPACITY);
        match remote {
            None => {}
            Some(remote) => {
                remotes_by_unique.insert(remote);
                remotes_by_handle.push(remote);
            }
        }
        storage.remotes_by_unique.insert(job.unique.clone(), remotes_by_unique);
        storage.remotes_by_handle.insert(job.handle.clone(), remotes_by_handle);
    }

    fn get_job(&mut self, worker: &mut Worker) -> Option<Arc<Job>> {
        let mut storage = self.lock().unwrap();
        let mut job: Option<Arc<Job>> = None;
        debug!("{:?}", &worker);
        for func in worker.iter() {
            debug!("func = {:?}", &func);
            match storage.queues.get_mut(&func) {
                None => {}
                Some(prios) => {
                    debug!("found func with {} priority queues", prios.len());
                    let mut i = 0;
                    for q in prios {
                        debug!("searching priority {}", i);
                        i = i + 1;
                        if !q.is_empty() {
                            debug!("Queue has items! {:?}", q.len());
                            job = q.pop_front();
                            break;
                        }
                    }
                }
            }
        }
        match job {
            Some(job) => {
                worker.assign_job(&job);
                Some(job)
            }
            None => None,
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation for SGX using RDRAND instruction
use crate::Error;
use core::mem;
use core::arch::x86_64::{__cpuid, _rdrand64_step};
use core::num::NonZeroU32;
use lazy_static::lazy_static;

// Recommendation from "Intel® Digital Random Number Generator (DRNG) Software
// Implementation Guide" - Section 5.2.1 and "Intel® 64 and IA-32 Architectures
// Software Developer’s Manual" - Volume 1 - Section 7.3.17.1.
const RETRY_LIMIT: usize = 10;
const WORD_SIZE: usize = mem::size_of::<u64>();

#[target_feature(enable = "rdrand")]
unsafe fn rdrand() -> Result<[u8; WORD_SIZE], Error> {
    for _ in 0..RETRY_LIMIT {
        let mut el = mem::uninitialized();
        if _rdrand64_step(&mut el) == 1 {
            return Ok(el.to_ne_bytes());
        }
    }
    error!("RDRAND failed, CPU issue likely");
    Err(Error::UNKNOWN)
}

// TODO use is_x86_feature_detected!("rdrand") when that works in core. See:
//   https://github.com/rust-lang-nursery/stdsimd/issues/464
fn is_rdrand_supported() -> bool {
    if cfg!(target_feature = "rdrand") {
        true
    } else if cfg!(target_env = "sgx") {
        false // No CPUID in SGX enclaves
    } else {
        // SAFETY: All x86_64 CPUs support CPUID leaf 1
        const FLAG: u32 = 1 << 30;
        lazy_static! {
            static ref HAS_RDRAND: bool = unsafe { __cpuid(1).ecx & FLAG != 0 };
        }
        *HAS_RDRAND
    }
}

pub fn getrandom_inner(dest: &mut [u8]) -> Result<(), Error> {
    if !is_rdrand_supported() {
        return Err(Error::UNAVAILABLE);
    }

    // SAFETY: After this point, rdrand is supported, so calling the rdrand
    // functions is not undefined behavior.
    unsafe { rdrand_exact(dest) }
}

#[target_feature(enable = "rdrand")]
unsafe fn rdrand_exact(dest: &mut [u8]) -> Result<(), Error> {
    // We use chunks_exact_mut instead of chunks_mut as it allows almost all
    // calls to memcpy to be elided by the compiler.
    let mut chunks = dest.chunks_exact_mut(WORD_SIZE);
    for chunk in chunks.by_ref() {
        chunk.copy_from_slice(&rdrand()?);
    }

    let tail = chunks.into_remainder();
    let n = tail.len();
    if n > 0 {
        tail.copy_from_slice(&rdrand()?[..n]);
    }
    Ok(())
}

#[inline(always)]
pub fn error_msg_inner(_: NonZeroU32) -> Option<&'static str> { None }
<commit_msg>Add error for sgx && !rdrand<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation for SGX using RDRAND instruction
use crate::Error;
use core::mem;
use core::arch::x86_64::{__cpuid, _rdrand64_step};
use core::num::NonZeroU32;
use lazy_static::lazy_static;

// Recommendation from "Intel® Digital Random Number Generator (DRNG) Software
// Implementation Guide" - Section 5.2.1 and "Intel® 64 and IA-32 Architectures
// Software Developer’s Manual" - Volume 1 - Section 7.3.17.1.
const RETRY_LIMIT: usize = 10;
const WORD_SIZE: usize = mem::size_of::<u64>();

#[target_feature(enable = "rdrand")]
unsafe fn rdrand() -> Result<[u8; WORD_SIZE], Error> {
    for _ in 0..RETRY_LIMIT {
        let mut el = mem::uninitialized();
        if _rdrand64_step(&mut el) == 1 {
            return Ok(el.to_ne_bytes());
        }
    }
    error!("RDRAND failed, CPU issue likely");
    Err(Error::UNKNOWN)
}

#[cfg(and(target_env = "sgx", not(target_feature = "rdrand")))]
compile_error!("SGX targets must enable RDRAND to get randomness");

// TODO use is_x86_feature_detected!("rdrand") when that works in core. See:
//   https://github.com/rust-lang-nursery/stdsimd/issues/464
fn is_rdrand_supported() -> bool {
    if cfg!(target_feature = "rdrand") {
        true
    } else {
        // SAFETY: All x86_64 CPUs support CPUID leaf 1
        const FLAG: u32 = 1 << 30;
        lazy_static! {
            static ref HAS_RDRAND: bool = unsafe { __cpuid(1).ecx & FLAG != 0 };
        }
        *HAS_RDRAND
    }
}

pub fn getrandom_inner(dest: &mut [u8]) -> Result<(), Error> {
    if !is_rdrand_supported() {
        return Err(Error::UNAVAILABLE);
    }

    // SAFETY: After this point, rdrand is supported, so calling the rdrand
    // functions is not undefined behavior.
    unsafe { rdrand_exact(dest) }
}

#[target_feature(enable = "rdrand")]
unsafe fn rdrand_exact(dest: &mut [u8]) -> Result<(), Error> {
    // We use chunks_exact_mut instead of chunks_mut as it allows almost all
    // calls to memcpy to be elided by the compiler.
    let mut chunks = dest.chunks_exact_mut(WORD_SIZE);
    for chunk in chunks.by_ref() {
        chunk.copy_from_slice(&rdrand()?);
    }

    let tail = chunks.into_remainder();
    let n = tail.len();
    if n > 0 {
        tail.copy_from_slice(&rdrand()?[..n]);
    }
    Ok(())
}

#[inline(always)]
pub fn error_msg_inner(_: NonZeroU32) -> Option<&'static str> { None }
<|endoftext|>


<commit_before>//! Implements the [`FrameRunner`](./runner/struct.FrameRunner.html) struct for managing frame simulations.
use chrono;

use clock::{GameTime, GameClock};
use framerate::counter::FrameCount;
use step::TimeStep;

/// A helper type for running frame simulations with a frame counter.
///
/// `FrameRunner` combines a [`GameClock`](../clock/struct.GameTime.html)
/// and a [`FrameCount`](../framerate/counter/struct.FrameCounter.html) object,
/// tracking each frame as well as a framerate. `FrameRunner` provides
/// a `tick` method like `GameClock`, and updates both the `GameClock` and
/// `FrameCount` objects contained.
#[derive(Debug)]
pub struct FrameRunner<C: FrameCount> {
    clock: GameClock,
    counter: C,
}

impl<C> FrameRunner<C>
where
    C: FrameCount,
{
    /// Construct a new `FrameRunner` from a `GameClock` and a `FrameCount`.
    pub fn new(clock: GameClock, counter: C) -> FrameRunner<C> {
        FrameRunner { clock, counter }
    }

    /// Get a reference to the contained `GameClock`.
    pub fn clock(&self) -> &GameClock {
        &self.clock
    }
    /// Get a mutable reference to the contained `GameClock`.
    pub fn clock_mut(&mut self) -> &mut GameClock {
        &mut self.clock
    }
    /// Get a reference to the contained `FrameCount`.
    pub fn counter(&self) -> &C {
        &self.counter
    }
    /// Get a mutable reference to the contained `FrameCount`.
    pub fn counter_mut(&mut self) -> &mut C {
        &mut self.counter
    }
    /// Mark the start of a new frame, updating time and frame rate statistics.
    ///
    /// The `GameTime` for the new frame is returned, with the same properties as that
    /// returned from [`GameClock::tick`](../clock/struct.GameClock.html#method.tick).
    pub fn tick<T: TimeStep>(&mut self, time_step: &T) -> GameTime {
        self.clock.tick(time_step)
    }

    /// Mark the start of a new frame with a specified wall time, updating time statistics.
    ///
    /// This function is like `tick` but allows for the start time for the
    /// frame to be specified.
    pub fn tick_with_wall_time<T: TimeStep>(
        &mut self,
        time_step: &T,
        frame_start: chrono::DateTime<chrono::Local>,
    ) -> GameTime {
        self.clock.tick_with_wall_time(time_step, frame_start)
    }

    /// Perform one frame of the simulation using `frame_fn`.
    ///
    /// The closure is passed the `GameTime` for the frame by calling `tick`
    /// and will call
    /// [`GameClock::sleep_remaining`](../clock/struct.GameClock.html#method.sleep_remaining)
    /// after the closure has ended.
    pub fn do_frame<T, F>(&mut self, time_step: &T, frame_fn: F)
    where
        T: TimeStep,
        F: FnOnce(GameTime),
    {
        let time = self.tick(time_step);
        frame_fn(time);
        self.clock.sleep_remaining(&self.counter);
    }
}
<commit_msg>FrameRunner now ticks the counter too.<commit_after>//! Implements the [`FrameRunner`](./runner/struct.FrameRunner.html) struct for managing frame simulations.
use chrono;

use clock::{GameTime, GameClock};
use framerate::counter::FrameCount;
use step::TimeStep;

/// A helper type for running frame simulations with a frame counter.
///
/// `FrameRunner` combines a [`GameClock`](../clock/struct.GameTime.html)
/// and a [`FrameCount`](../framerate/counter/struct.FrameCounter.html) object,
/// tracking each frame as well as a framerate. `FrameRunner` provides
/// a `tick` method like `GameClock`, and updates both the `GameClock` and
/// `FrameCount` objects contained.
#[derive(Debug)]
pub struct FrameRunner<C: FrameCount> {
    clock: GameClock,
    counter: C,
}

impl<C> FrameRunner<C>
where
    C: FrameCount,
{
    /// Construct a new `FrameRunner` from a `GameClock` and a `FrameCount`.
    pub fn new(clock: GameClock, counter: C) -> FrameRunner<C> {
        FrameRunner { clock, counter }
    }

    /// Get a reference to the contained `GameClock`.
    pub fn clock(&self) -> &GameClock {
        &self.clock
    }
    /// Get a mutable reference to the contained `GameClock`.
    pub fn clock_mut(&mut self) -> &mut GameClock {
        &mut self.clock
    }
    /// Get a reference to the contained `FrameCount`.
    pub fn counter(&self) -> &C {
        &self.counter
    }
    /// Get a mutable reference to the contained `FrameCount`.
    pub fn counter_mut(&mut self) -> &mut C {
        &mut self.counter
    }
    /// Mark the start of a new frame, updating time and frame rate statistics.
    ///
    /// The `GameTime` for the new frame is returned, with the same properties as that
    /// returned from [`GameClock::tick`](../clock/struct.GameClock.html#method.tick).
    pub fn tick<T: TimeStep>(&mut self, time_step: &T) -> GameTime {
        let time = self.clock.tick(time_step);
        self.counter.tick(&time);
        time
    }

    /// Mark the start of a new frame with a specified wall time, updating time statistics.
    ///
    /// This function is like `tick` but allows for the start time for the
    /// frame to be specified.
    pub fn tick_with_wall_time<T: TimeStep>(
        &mut self,
        time_step: &T,
        frame_start: chrono::DateTime<chrono::Local>,
    ) -> GameTime {
        let time = self.clock.tick_with_wall_time(time_step, frame_start);
        self.counter.tick(&time);
        time
    }

    /// Perform one frame of the simulation using `frame_fn`.
    ///
    /// The closure is passed the `GameTime` for the frame by calling `tick`
    /// and will call
    /// [`GameClock::sleep_remaining`](../clock/struct.GameClock.html#method.sleep_remaining)
    /// after the closure has ended.
    pub fn do_frame<T, F>(&mut self, time_step: &T, frame_fn: F)
    where
        T: TimeStep,
        F: FnOnce(GameTime),
    {
        let time = self.tick(time_step);
        frame_fn(time);
        self.clock.sleep_remaining(&self.counter);
    }
}
<|endoftext|>


<commit_before>//! Running of programs.

use parser::Instruction;
use std;

///
/// A runner of programs.
///
#[derive(Default)]
pub struct Runner {
    /// Program to be run.
    prog: Vec<Instruction>,
    /// A potentially infinite number of data cells.
    data: Vec<u8>,
    /// Instruction pointer (index).
    ip: usize,
    /// Data-cell pointer (index).
    dp: usize
}

impl Runner {
    ///
    /// Constructs a new runner.
    ///
    pub fn new() -> Self {
        Runner::default()
    }

    ///
    /// Runs the given program.
    ///
    /// # Arguments
    ///
    /// * `prog` - Program to be run.
    /// * `input` - Input for the program.
    /// * `output` - Output for the program.
    ///
    pub fn run(&mut self,
               prog: Vec<Instruction>,
               mut input: &mut std::io::Read,
               mut output: &mut std::io::Write) -> Result<(), String> {
        self.prog = prog;
        self.data.clear();
        self.ip = 0;
        self.dp = 0;

        while !self.is_prog_end() {
            try!(self.run_instruction(&mut input, &mut output));
        }

        Ok(())
    }

    fn is_prog_end(&self) -> bool {
        self.ip >= self.prog.len()
    }

    fn run_instruction(&mut self,
                       mut input: &mut std::io::Read,
                       mut output: &mut std::io::Write) -> Result<(), String> {
        match self.prog[self.ip] {
            Instruction::Next => {
                self.dp += 1;
            },
            Instruction::Prev => {
                self.dp -= 1;
            },
            Instruction::Inc => {
                let value = self.get_value();
                self.set_value(value.wrapping_add(1));
            },
            Instruction::Dec => {
                let value = self.get_value();
                self.set_value(value.wrapping_sub(1));
            },
            Instruction::Read => {
                let value = try!(self.read_value(&mut input));
                self.set_value(value);
            },
            Instruction::Write => {
                let value = self.get_value();
                try!(self.write_value(&mut output, value));
            },
            Instruction::LoopStart(loop_end_ip) => {
                if self.is_current_cell_zero() {
                    self.ip = loop_end_ip;
                }
            },
            Instruction::LoopEnd(loop_start_ip) => {
                if !self.is_current_cell_zero() {
                    self.ip = loop_start_ip;
                }
            }
        }
        self.ip += 1;
        Ok(())
    }

    fn get_value(&mut self) -> u8 {
        self.ensure_current_cell_is_accessible();
        self.data[self.dp]
    }

    fn set_value(&mut self, value: u8) {
        self.ensure_current_cell_is_accessible();
        self.data[self.dp] = value;
    }

    fn read_value(&self, mut input: &mut std::io::Read) -> Result<u8, String> {
        let mut buf = [0u8];
        match input.read_exact(&mut buf[..]) {
            Err(err) => Err(format!("{}", err)),
            Ok(_) => Ok(buf[0])
        }
    }

    fn write_value(&self, mut output: &mut std::io::Write, value: u8) -> Result<(), String> {
        match output.write(&[value]) {
            Err(err) => Err(format!("{}", err)),
            Ok(_) => Ok(())
        }
    }

    fn is_current_cell_zero(&mut self) -> bool {
        self.get_value() == 0
    }

    fn ensure_current_cell_is_accessible(&mut self) {
        let cell_count = self.data.len();
        if self.dp >= cell_count {
            for _ in 0..(self.dp - cell_count + 1) {
                self.data.push(0);
            }
        }
    }
}

///
/// Runs the given program.
///
/// # Arguments
///
/// * `prog` - Program to be run.
/// * `input` - Input for the program.
/// * `output` - Output for the program.
///
pub fn run(prog: Vec<Instruction>,
       mut input: &mut std::io::Read,
       mut output: &mut std::io::Write) -> Result<(), String> {
    let mut runner = Runner::new();
    runner.run(prog, &mut input, &mut output)
}

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

    fn assert_run_writes_correct_output(prog: Vec<Instruction>,
                                               input: &[u8],
                                               expected_output: &[u8]) {
        let mut output = Vec::new();
        {
            let mut input_reader = std::io::BufReader::new(input.as_ref());
            let mut output_writer = std::io::BufWriter::new(&mut output);
            run(prog, &mut input_reader, &mut output_writer).unwrap();
        }

        assert_eq!(output, expected_output);
    }

    #[test]
    fn test_run_does_nothing_when_program_has_no_instructions() {
        assert_run_writes_correct_output(
            vec![],
            &[],
            &[]
        );
    }

    #[test]
    fn test_run_can_read_from_input_and_write_to_output() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::Write
            ],
            &[1],
            &[1]
        );
    }

    #[test]
    fn test_run_can_increment_cell() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::Inc,
                Instruction::Write,
            ],
            &[1],
            &[2]
        );
    }

    #[test]
    fn test_run_can_decrement_cell() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::Dec,
                Instruction::Write,
            ],
            &[1],
            &[0]
        );
    }

    #[test]
    fn test_run_can_move_between_cells() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::Next,
                Instruction::Read,
                Instruction::Next,
                Instruction::Prev,
                Instruction::Write,
                Instruction::Prev,
                Instruction::Write,
            ],
            &[1, 2],
            &[2, 1]
        );
    }

    #[test]
    fn test_run_can_loop() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::LoopStart(6),
                Instruction::Dec,
                Instruction::Next,
                Instruction::Inc,
                Instruction::Prev,
                Instruction::LoopEnd(1),
                Instruction::Write,
                Instruction::Next,
                Instruction::Write
            ],
            &[2],
            &[0, 2]
        );
    }
}
<commit_msg>Rename get_value() to load_value() and set_value() to store_value().<commit_after>//! Running of programs.

use parser::Instruction;
use std;

///
/// A runner of programs.
///
#[derive(Default)]
pub struct Runner {
    /// Program to be run.
    prog: Vec<Instruction>,
    /// A potentially infinite number of data cells.
    data: Vec<u8>,
    /// Instruction pointer (index).
    ip: usize,
    /// Data-cell pointer (index).
    dp: usize
}

impl Runner {
    ///
    /// Constructs a new runner.
    ///
    pub fn new() -> Self {
        Runner::default()
    }

    ///
    /// Runs the given program.
    ///
    /// # Arguments
    ///
    /// * `prog` - Program to be run.
    /// * `input` - Input for the program.
    /// * `output` - Output for the program.
    ///
    pub fn run(&mut self,
               prog: Vec<Instruction>,
               mut input: &mut std::io::Read,
               mut output: &mut std::io::Write) -> Result<(), String> {
        self.prog = prog;
        self.data.clear();
        self.ip = 0;
        self.dp = 0;

        while !self.is_prog_end() {
            try!(self.run_instruction(&mut input, &mut output));
        }

        Ok(())
    }

    fn is_prog_end(&self) -> bool {
        self.ip >= self.prog.len()
    }

    fn run_instruction(&mut self,
                       mut input: &mut std::io::Read,
                       mut output: &mut std::io::Write) -> Result<(), String> {
        match self.prog[self.ip] {
            Instruction::Next => {
                self.dp += 1;
            },
            Instruction::Prev => {
                self.dp -= 1;
            },
            Instruction::Inc => {
                let value = self.load_value();
                self.store_value(value.wrapping_add(1));
            },
            Instruction::Dec => {
                let value = self.load_value();
                self.store_value(value.wrapping_sub(1));
            },
            Instruction::Read => {
                let value = try!(self.read_value(&mut input));
                self.store_value(value);
            },
            Instruction::Write => {
                let value = self.load_value();
                try!(self.write_value(&mut output, value));
            },
            Instruction::LoopStart(loop_end_ip) => {
                if self.is_current_cell_zero() {
                    self.ip = loop_end_ip;
                }
            },
            Instruction::LoopEnd(loop_start_ip) => {
                if !self.is_current_cell_zero() {
                    self.ip = loop_start_ip;
                }
            }
        }
        self.ip += 1;
        Ok(())
    }

    fn load_value(&mut self) -> u8 {
        self.ensure_current_cell_is_accessible();
        self.data[self.dp]
    }

    fn store_value(&mut self, value: u8) {
        self.ensure_current_cell_is_accessible();
        self.data[self.dp] = value;
    }

    fn read_value(&self, mut input: &mut std::io::Read) -> Result<u8, String> {
        let mut buf = [0u8];
        match input.read_exact(&mut buf[..]) {
            Err(err) => Err(format!("{}", err)),
            Ok(_) => Ok(buf[0])
        }
    }

    fn write_value(&self, mut output: &mut std::io::Write, value: u8) -> Result<(), String> {
        match output.write(&[value]) {
            Err(err) => Err(format!("{}", err)),
            Ok(_) => Ok(())
        }
    }

    fn is_current_cell_zero(&mut self) -> bool {
        self.load_value() == 0
    }

    fn ensure_current_cell_is_accessible(&mut self) {
        let cell_count = self.data.len();
        if self.dp >= cell_count {
            for _ in 0..(self.dp - cell_count + 1) {
                self.data.push(0);
            }
        }
    }
}

///
/// Runs the given program.
///
/// # Arguments
///
/// * `prog` - Program to be run.
/// * `input` - Input for the program.
/// * `output` - Output for the program.
///
pub fn run(prog: Vec<Instruction>,
       mut input: &mut std::io::Read,
       mut output: &mut std::io::Write) -> Result<(), String> {
    let mut runner = Runner::new();
    runner.run(prog, &mut input, &mut output)
}

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

    fn assert_run_writes_correct_output(prog: Vec<Instruction>,
                                               input: &[u8],
                                               expected_output: &[u8]) {
        let mut output = Vec::new();
        {
            let mut input_reader = std::io::BufReader::new(input.as_ref());
            let mut output_writer = std::io::BufWriter::new(&mut output);
            run(prog, &mut input_reader, &mut output_writer).unwrap();
        }

        assert_eq!(output, expected_output);
    }

    #[test]
    fn test_run_does_nothing_when_program_has_no_instructions() {
        assert_run_writes_correct_output(
            vec![],
            &[],
            &[]
        );
    }

    #[test]
    fn test_run_can_read_from_input_and_write_to_output() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::Write
            ],
            &[1],
            &[1]
        );
    }

    #[test]
    fn test_run_can_increment_cell() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::Inc,
                Instruction::Write,
            ],
            &[1],
            &[2]
        );
    }

    #[test]
    fn test_run_can_decrement_cell() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::Dec,
                Instruction::Write,
            ],
            &[1],
            &[0]
        );
    }

    #[test]
    fn test_run_can_move_between_cells() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::Next,
                Instruction::Read,
                Instruction::Next,
                Instruction::Prev,
                Instruction::Write,
                Instruction::Prev,
                Instruction::Write,
            ],
            &[1, 2],
            &[2, 1]
        );
    }

    #[test]
    fn test_run_can_loop() {
        assert_run_writes_correct_output(
            vec![
                Instruction::Read,
                Instruction::LoopStart(6),
                Instruction::Dec,
                Instruction::Next,
                Instruction::Inc,
                Instruction::Prev,
                Instruction::LoopEnd(1),
                Instruction::Write,
                Instruction::Next,
                Instruction::Write
            ],
            &[2],
            &[0, 2]
        );
    }
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate rustc_serialize;
extern crate rusqlite;
extern crate jmap;
extern crate time;

#[macro_use]
extern crate log;

mod logger;
mod db;
mod util;
mod record;

use std::default::Default;
use std::io::{Read, Write, ErrorKind};

use std::fs::File;

use hyper::server::{Request, Response};
use hyper::method::Method::{Post, Get, Head};
use hyper::status::StatusCode;
use hyper::uri::RequestUri::AbsolutePath;
use hyper::header;

use rustc_serialize::json::{Json,ToJson};

use jmap::parse::FromJson;
use jmap::method::{RequestBatch, ResponseBatch, ResponseMethod};
use jmap::method::RequestMethod::*;
use jmap::method::ResponseMethod::*;

use util::RequestContext;
use record::RecordHandler;
use db::Db;

struct StatusBody {
    code: StatusCode,
    body: Option<Vec<u8>>
}
impl StatusBody {
    fn new(code: StatusCode, body: Option<Vec<u8>>) -> StatusBody {
        StatusBody { code: code, body: body }
    }
}

macro_rules! make_crud_method_dispatcher {
    ($method: expr, $rmethods: expr, $r: expr,
     $($ty: ty { $get: ident => $rget: ident, $set: ident => $rset: ident, $getup: ident => $rgetup: ident }),*) => {
        $rmethods = match $method {
            $(
                $get(args, id)
                    => vec!($r.get_records(&args).map(|a| $rget(a, id.clone())).unwrap_or_else(|e| ResponseError(e, id.clone()))),
                $set(args, id)
                    => vec!($r.set_records(&args).map(|a| $rset(a, id.clone())).unwrap_or_else(|e| ResponseError(e, id.clone()))),
                $getup(args, id) => $r.get_record_updates(&args).map(|a| {
                    match a {
                        (a, Some(b)) => vec!($rgetup(a, id.clone()), $rget(b, id.clone())),
                        (a, _)       => vec!($rgetup(a, id.clone())),
                    }
                }).unwrap_or_else(|e| vec!(ResponseError(e, id.clone()))),
            )*
            RequestError(args, id) => vec!(ResponseError(args, id)),
        };
    }
}

fn jmap_handler(mut req: Request) -> StatusBody {

    match Json::from_reader(&mut req) {
        Ok(j) => match RequestBatch::from_json(&j) {
            Ok(b) => {
                let mut rbatch: ResponseBatch = ResponseBatch::default();

                let r = RequestContext {
                    userid: 1, // XXX get userid from auth
                    db: Db::open().unwrap(),
                };

                for method in b.0.into_iter() {
                    let rmethods: Vec<ResponseMethod>;
                    make_crud_method_dispatcher!(method, rmethods, r,
                        Calendar {
                            GetCalendars       => Calendars,
                            SetCalendars       => CalendarsSet,
                            GetCalendarUpdates => CalendarUpdates
                        },
                        CalendarEvent {
                            GetCalendarEvents       => CalendarEvents,
                            SetCalendarEvents       => CalendarEventsSet,
                            GetCalendarEventUpdates => CalendarEventUpdates
                        },
                        Contact {
                            GetContacts       => Contacts,
                            SetContacts       => ContactsSet,
                            GetContactUpdates => ContactUpdates
                        },
                        ContactGroup {
                            GetContactGroups       => ContactGroups,
                            SetContactGroups       => ContactGroupsSet,
                            GetContactGroupUpdates => ContactGroupUpdates
                        },
                        Mailbox {
                            GetMailboxes      => Mailboxes,
                            SetMailboxes      => MailboxesSet,
                            GetMailboxUpdates => MailboxUpdates
                        }
                    );

                    rbatch.0.extend(rmethods.into_iter());
                }

                StatusBody::new(StatusCode::Ok, Some(rbatch.to_json().to_string().into_bytes()))
            },
            Err(e) =>
                StatusBody::new(StatusCode::BadRequest, Some(e.to_string().into_bytes())),
        },
        Err(e) => StatusBody::new(StatusCode::BadRequest, Some(e.to_string().into_bytes())),
    }
}

fn file_handler(path: &String, include_body: bool) -> StatusBody {
    match File::open(String::from("client") + if path == "/" { "/index.html" } else { path }) {
        Err(ref e) => match e.kind() {
            ErrorKind::NotFound => StatusBody::new(StatusCode::NotFound, None),
            // XXX others
            _ => StatusBody::new(StatusCode::InternalServerError, Some(format!("{}", e).into_bytes())),
        },
        Ok(ref mut f) => {
            match include_body {
                false => StatusBody::new(StatusCode::Ok, None),
                true  => {
                    let mut s = String::new();
                    match f.read_to_string(&mut s) {
                        Err(ref e) =>
                            StatusBody::new(StatusCode::InternalServerError, Some(format!("{}", e).into_bytes())),
                        _ =>
                            StatusBody::new(StatusCode::Ok, Some(s.into_bytes())),
                    }
                },
            }
        },
    }
}

fn finish_response(mut res: Response, out: StatusBody) {
    *res.status_mut() = out.code;

    match res.start()
        .and_then(|mut res|
                  match out.body {
                      Some(ref b) => {
                          try!(res.write_all(b));
                          Ok(res)
                      }
                      None => {
                          Ok(res)
                      }
                  })
        .and_then(|res| res.end()) {
            Err(e) => error!("response error: {}", e),
            _      => (),
        };
}

fn http_handler(req: Request, mut res: Response) {
    res.headers_mut().set(header::Server("salada/0.0.5".to_string()));

    let method = req.method.clone();
    let uri = req.uri.clone();

    match (method, uri) {
        (Post, AbsolutePath(ref path)) if path == "/jmap" =>
            finish_response(res, jmap_handler(req)),

        (Get, AbsolutePath(ref path)) =>
            finish_response(res, file_handler(path, true)),

        (Head, AbsolutePath(ref path)) =>
            finish_response(res, file_handler(path, false)),

        _ => finish_response(res, StatusBody::new(StatusCode::MethodNotAllowed, None)),
    };
}

fn main() {
    logger::init().unwrap();

    info!("Listening on http://127.0.0.1:3000/jmap");
    hyper::Server::http("127.0.0.1:3000").unwrap().handle(http_handler).unwrap();
}
<commit_msg>Add request logging, drain request on error<commit_after>extern crate hyper;
extern crate rustc_serialize;
extern crate rusqlite;
extern crate jmap;
extern crate time;

#[macro_use]
extern crate log;

mod logger;
mod db;
mod util;
mod record;

use std::default::Default;
use std::io::{Read, Write, ErrorKind};

use std::fs::File;

use hyper::server::{Request, Response};
use hyper::method::Method;
use hyper::method::Method::{Post, Get, Head};
use hyper::status::StatusCode;
use hyper::uri::RequestUri::AbsolutePath;
use hyper::header;

use rustc_serialize::json::{Json,ToJson};

use jmap::parse::FromJson;
use jmap::method::{RequestBatch, ResponseBatch, ResponseMethod};
use jmap::method::RequestMethod::*;
use jmap::method::ResponseMethod::*;

use util::RequestContext;
use record::RecordHandler;
use db::Db;

struct StatusBody {
    code: StatusCode,
    body: Option<Vec<u8>>
}
impl StatusBody {
    fn new(code: StatusCode, body: Option<Vec<u8>>) -> StatusBody {
        StatusBody { code: code, body: body }
    }
}

macro_rules! make_crud_method_dispatcher {
    ($method: expr, $rmethods: expr, $r: expr,
     $($ty: ty { $get: ident => $rget: ident, $set: ident => $rset: ident, $getup: ident => $rgetup: ident }),*) => {
        $rmethods = match $method {
            $(
                $get(args, id)
                    => vec!($r.get_records(&args).map(|a| $rget(a, id.clone())).unwrap_or_else(|e| ResponseError(e, id.clone()))),
                $set(args, id)
                    => vec!($r.set_records(&args).map(|a| $rset(a, id.clone())).unwrap_or_else(|e| ResponseError(e, id.clone()))),
                $getup(args, id) => $r.get_record_updates(&args).map(|a| {
                    match a {
                        (a, Some(b)) => vec!($rgetup(a, id.clone()), $rget(b, id.clone())),
                        (a, _)       => vec!($rgetup(a, id.clone())),
                    }
                }).unwrap_or_else(|e| vec!(ResponseError(e, id.clone()))),
            )*
            RequestError(args, id) => vec!(ResponseError(args, id)),
        };
    }
}

fn jmap_handler(mut req: Request) -> StatusBody {

    match Json::from_reader(&mut req) {
        Ok(j) => match RequestBatch::from_json(&j) {
            Ok(b) => {
                let mut rbatch: ResponseBatch = ResponseBatch::default();

                let r = RequestContext {
                    userid: 1, // XXX get userid from auth
                    db: Db::open().unwrap(),
                };

                for method in b.0.into_iter() {
                    let rmethods: Vec<ResponseMethod>;
                    make_crud_method_dispatcher!(method, rmethods, r,
                        Calendar {
                            GetCalendars       => Calendars,
                            SetCalendars       => CalendarsSet,
                            GetCalendarUpdates => CalendarUpdates
                        },
                        CalendarEvent {
                            GetCalendarEvents       => CalendarEvents,
                            SetCalendarEvents       => CalendarEventsSet,
                            GetCalendarEventUpdates => CalendarEventUpdates
                        },
                        Contact {
                            GetContacts       => Contacts,
                            SetContacts       => ContactsSet,
                            GetContactUpdates => ContactUpdates
                        },
                        ContactGroup {
                            GetContactGroups       => ContactGroups,
                            SetContactGroups       => ContactGroupsSet,
                            GetContactGroupUpdates => ContactGroupUpdates
                        },
                        Mailbox {
                            GetMailboxes      => Mailboxes,
                            SetMailboxes      => MailboxesSet,
                            GetMailboxUpdates => MailboxUpdates
                        }
                    );

                    rbatch.0.extend(rmethods.into_iter());
                }

                StatusBody::new(StatusCode::Ok, Some(rbatch.to_json().to_string().into_bytes()))
            },
            Err(e) =>
                StatusBody::new(StatusCode::BadRequest, Some(e.to_string().into_bytes())),
        },
        Err(e) => StatusBody::new(StatusCode::BadRequest, Some(e.to_string().into_bytes())),
    }
}

fn file_handler(path: &String, include_body: bool) -> StatusBody {
    match File::open(String::from("client") + if path == "/" { "/index.html" } else { path }) {
        Err(ref e) => match e.kind() {
            ErrorKind::NotFound => StatusBody::new(StatusCode::NotFound, None),
            // XXX others
            _ => StatusBody::new(StatusCode::InternalServerError, Some(format!("{}", e).into_bytes())),
        },
        Ok(ref mut f) => {
            match include_body {
                false => StatusBody::new(StatusCode::Ok, None),
                true  => {
                    let mut s = String::new();
                    match f.read_to_string(&mut s) {
                        Err(ref e) =>
                            StatusBody::new(StatusCode::InternalServerError, Some(format!("{}", e).into_bytes())),
                        _ =>
                            StatusBody::new(StatusCode::Ok, Some(s.into_bytes())),
                    }
                },
            }
        },
    }
}

fn finish_response(method: Method, path: &String, mut res: Response, out: StatusBody) {
    *res.status_mut() = out.code;

    info!("HTTP: {} {} => {}", method, path, out.code);

    match res.start()
        .and_then(|mut res|
                  match out.body {
                      Some(ref b) => {
                          try!(res.write_all(b));
                          Ok(res)
                      }
                      None => {
                          Ok(res)
                      }
                  })
        .and_then(|res| res.end()) {
            Err(e) => error!("response error: {}", e),
            _      => (),
        };
}

fn http_handler(mut req: Request, mut res: Response) {
    res.headers_mut().set(header::Server("salada/0.0.5".to_string()));

    let method = req.method.clone();
    let uri = req.uri.clone();

    match (method, uri) {
        (Post, AbsolutePath(ref path)) if path == "/jmap" =>
            finish_response(Post, path, res, jmap_handler(req)),

        (Get, AbsolutePath(ref path)) =>
            finish_response(Get, path, res, file_handler(path, true)),

        (Head, AbsolutePath(ref path)) =>
            finish_response(Head, path, res, file_handler(path, false)),

        (_, ref ap) => {
            let s = match ap {
                &AbsolutePath(ref u) => u,
                _ => panic!("RequestUri not AbsolutePath?"),
            };
            let mut drain: Vec<u8> = vec!();
            req.read_to_end(&mut drain).ok();
            finish_response(req.method, &s, res, StatusBody::new(StatusCode::MethodNotAllowed, None))
        },
    };
}

fn main() {
    logger::init().unwrap();

    info!("Listening on http://127.0.0.1:3000/jmap");
    hyper::Server::http("127.0.0.1:3000").unwrap().handle(http_handler).unwrap();
}
<|endoftext|>