<commit_before>// Copyright (c) 2014 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Miscellaneous utilities for low level networking

use bindings::libc;

use std::fmt;
use std::str::FromStr;
use std::mem;
use std::num::from_str_radix;
use std::io::net::ip::IpAddr;

#[cfg(not(windows))] use internal;

/// A MAC address
#[deriving(PartialEq, Eq, Clone)]
pub struct MacAddr(pub u8, pub u8, pub u8, pub u8, pub u8, pub u8);
impl Copy for MacAddr {}

impl fmt::Show for MacAddr {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MacAddr(a, b, c, d, e, f) =>
                write!(fmt, "{:x}:{:x}:{:x}:{:x}:{:x}:{:x}",
                       a, b, c, d, e, f)
        }
    }
}

impl FromStr for MacAddr {
    fn from_str(s: &str) -> Option<MacAddr> {
        let mut parts = [0u8, ..6];
        let mut splits = s.split(':');
        let mut i = 0;
        for split in splits {
            if i == 6 {
                return None;
            }
            match from_str_radix(split, 16) {
                Some(b) if split.len() != 0 => parts[i] = b,
                _ => return None,
            }
            i += 1;
        }

        if i == 6 {
            Some(MacAddr(parts[0], parts[1], parts[2], parts[3], parts[4], parts[5]))
        } else {
            None
        }
    }
}

#[test]
fn mac_addr_from_str() {
    assert_eq!(from_str::<MacAddr>("00:00:00:00:00:00"), Some(MacAddr(0, 0, 0, 0, 0, 0)));
    assert_eq!(from_str::<MacAddr>("ff:ff:ff:ff:ff:ff"), Some(MacAddr(0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF)));
    assert_eq!(from_str::<MacAddr>("12:34:56:78:90:ab"), Some(MacAddr(0x12, 0x34, 0x56, 0x78, 0x90, 0xAB)));
    assert_eq!(from_str::<MacAddr>("::::::"), None);
    assert_eq!(from_str::<MacAddr>("0::::::"), None);
    assert_eq!(from_str::<MacAddr>("::::0::"), None);
    assert_eq!(from_str::<MacAddr>("12:34:56:78"), None);
    assert_eq!(from_str::<MacAddr>("12:34:56:78:"), None);
    assert_eq!(from_str::<MacAddr>("12:34:56:78:90"), None);
    assert_eq!(from_str::<MacAddr>("12:34:56:78:90:"), None);
    assert_eq!(from_str::<MacAddr>("12:34:56:78:90:00:00"), None);
    assert_eq!(from_str::<MacAddr>("xx:xx:xx:xx:xx:xx"), None);
}

/// Represents a network interface and its associated addresses
#[deriving(Clone, PartialEq, Eq, Show)]
pub struct NetworkInterface {
    /// The name of the interface
    pub name: String,
    /// The interface index (operating system specific)
    pub index: u32,
    /// A MAC address for the interface
    pub mac: Option<MacAddr>,
    /// An IP addresses for the interface
    pub ips: Option<Vec<IpAddr>>,
    /// Operating system specific flags for the interface
    pub flags: u32,
}

impl NetworkInterface {
    /// Retreive the MAC address associated with the interface
    pub fn mac_address(&self) -> MacAddr {
        self.mac.unwrap()
    }

    /// Is the interface a loopback interface?
    pub fn is_loopback(&self) -> bool {
        self.flags & (libc::IFF_LOOPBACK as u32) != 0
    }
}

#[cfg(target_os = "linux")]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == libc::AF_PACKET {
            let sll: *const libc::sockaddr_ll = mem::transmute(sa);
            let mac = MacAddr((*sll).sll_addr[0], (*sll).sll_addr[1],
                              (*sll).sll_addr[2], (*sll).sll_addr[3],
                              (*sll).sll_addr[4], (*sll).sll_addr[5]);
            return (Some(mac), None);
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                        mem::size_of::<libc::sockaddr_storage>());
            return match addr {
                Ok(sa) => (None, Some(sa.ip)),
                Err(_) => (None, None)
            };
        }
    }
}

#[cfg(any(target_os = "freebsd", target_os = "macos"))]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use bindings::bpf;
    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == bpf::AF_LINK {
            let sdl: *const bpf::sockaddr_dl = mem::transmute(sa);
            let nlen = (*sdl).sdl_nlen as uint;
            let mac = MacAddr((*sdl).sdl_data[nlen + 0] as u8,
                              (*sdl).sdl_data[nlen + 1] as u8,
                              (*sdl).sdl_data[nlen + 2] as u8,
                              (*sdl).sdl_data[nlen + 3] as u8,
                              (*sdl).sdl_data[nlen + 4] as u8,
                              (*sdl).sdl_data[nlen + 5] as u8
                      );
            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                        mem::size_of::<libc::sockaddr_storage>());
            match addr {
                Ok(sa) => (None, Some(sa.ip)),
                Err(_) => (None, None)
            }
        }
    }
}

/// Get a list of available network interfaces for the current machine.
#[inline]
pub fn get_network_interfaces() -> Vec<NetworkInterface> {
    get_network_interfaces_impl()
}

#[cfg(not(windows))]
fn get_network_interfaces_impl() -> Vec<NetworkInterface> {
    use std::string::String;

    let mut ifaces: Vec<NetworkInterface> = Vec::new();
    unsafe {
        let mut addrs: *mut libc::ifaddrs = mem::uninitialized();
        if libc::getifaddrs(&mut addrs) != 0 {
            return ifaces;
        }
        let mut addr = addrs;
        while addr.is_not_null() {
            let name = String::from_raw_buf((*addr).ifa_name as *const u8);
            let (mac, ip) = sockaddr_to_network_addr((*addr).ifa_addr as *const libc::sockaddr);
            let ni = NetworkInterface {
                name: name.clone(),
                index: 0,
                mac: mac,
                ips: ip.map(|ip| [ip].to_vec()),
                flags: (*addr).ifa_flags
            };
            let mut found: bool = false;
            for iface in ifaces.iter_mut() {
                if name == iface.name {
                    merge(iface, &ni);
                    found = true;
                }
            }
            if !found {
                ifaces.push(ni);
            }

            addr = (*addr).ifa_next;
        }
        libc::freeifaddrs(addrs);

        for iface in ifaces.iter_mut() {
            iface.index = iface.name.with_c_str(
                |name| libc::if_nametoindex(name)
            );
        }
        return ifaces;
    }

    fn merge(old: &mut NetworkInterface, new: &NetworkInterface) {
        old.mac = match new.mac {
            None => old.mac,
            _ => new.mac
        };
        match (&mut old.ips, &new.ips) {
            (&Some(ref mut old_ips), &Some(ref new_ips)) => old_ips.push_all(new_ips.as_slice()),
            _ => {}
        };
        old.flags = old.flags | new.flags;
    }

}

#[cfg(windows)]
fn get_network_interfaces_impl() -> Vec<NetworkInterface> {
    use std::str::from_utf8;
    use std::string::raw;

    use bindings::winpcap;

    let mut adapters_size = 0u32;

    unsafe {
        let mut tmp: winpcap::IP_ADAPTER_INFO = mem::zeroed();
        // FIXME [windows] This only gets IPv4 addresses - should use GetAdaptersAddresses
        winpcap::GetAdaptersInfo(
            &mut tmp,
            &mut adapters_size
        );
    }

    let vec_size = adapters_size / mem::size_of::<winpcap::IP_ADAPTER_INFO>() as u32;

    let mut adapters = Vec::with_capacity(vec_size as uint);

    // FIXME [windows] Check return code
    unsafe {
        winpcap::GetAdaptersInfo(adapters.as_mut_ptr(), &mut adapters_size);
    }

    // Create a complete list of NetworkInterfaces for the machine
    let mut cursor = adapters.as_mut_ptr();
    let mut all_ifaces = Vec::with_capacity(vec_size as uint);
    while cursor.is_not_null() {
        let mac = unsafe {
                    MacAddr((*cursor).Address[0],
                            (*cursor).Address[1],
                            (*cursor).Address[2],
                            (*cursor).Address[3],
                            (*cursor).Address[4],
                            (*cursor).Address[5])
                  };
        let mut ip_cursor = unsafe { &mut (*cursor).IpAddressList as winpcap::PIP_ADDR_STRING};
        let mut ips: Vec<IpAddr> = Vec::new();
        while ip_cursor.is_not_null() {
            let ip_str = unsafe {
                            raw::from_buf((*ip_cursor).IpAddress.String.as_ptr() as *const u8)
                         };
            ips.push(from_str(ip_str.as_slice()).unwrap());
            ip_cursor = unsafe { (*ip_cursor).Next };
        }
        unsafe {
            all_ifaces.push(NetworkInterface {
                        name: raw::from_buf((*cursor).AdapterName.as_ptr() as *const u8),
                        index: (*cursor).Index,
                        mac: Some(mac),
                        ips: Some(ips),
                        //flags: (*cursor).Type, // FIXME [windows]
                        flags: 0,
                     });

            cursor = (*cursor).Next;
        }
    }

    let mut buf = [0u8, ..4096];
    let mut buflen = buf.len() as u32;

    // Gets list of supported adapters in form:
    // adapter1\0adapter2\0\0desc1\0desc2\0\0
    if unsafe { winpcap::PacketGetAdapterNames(buf.as_mut_ptr() as *mut i8, &mut buflen) } == 0 {
        // FIXME [windows] Should allocate a buffer big enough and try again
        //        - size should be buf.len() + buflen (buflen is overwritten)
        panic!("FIXME [windows] unable to get interface list");
    }

    let buf_str = from_utf8(buf).unwrap();
    let iface_names = buf_str.split_str("\0\0").next();
    let mut vec = Vec::new();

    // Return only supported adapters
    match iface_names {
        Some(iface_names) => {
            for iface in iface_names.split('\0') {
                let name = iface.to_string();
                let next = all_ifaces.iter()
                                     .filter(|x| name.as_slice().ends_with(x.name.as_slice()))
                                     .next();
                if next.is_some() {
                    let mut iface = next.unwrap().clone();
                    iface.name = name;
                    vec.push(iface);
                }
            }
        },
        None => ()
    };

    vec
}

<commit_msg>Fixes for from_str.<commit_after>// Copyright (c) 2014 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Miscellaneous utilities for low level networking

use bindings::libc;

use std::fmt;
use std::str::FromStr;
use std::mem;
use std::num::from_str_radix;
use std::io::net::ip::IpAddr;

#[cfg(not(windows))] use internal;

/// A MAC address
#[deriving(PartialEq, Eq, Clone)]
pub struct MacAddr(pub u8, pub u8, pub u8, pub u8, pub u8, pub u8);
impl Copy for MacAddr {}

impl fmt::Show for MacAddr {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MacAddr(a, b, c, d, e, f) =>
                write!(fmt, "{:x}:{:x}:{:x}:{:x}:{:x}:{:x}",
                       a, b, c, d, e, f)
        }
    }
}

impl FromStr for MacAddr {
    fn from_str(s: &str) -> Option<MacAddr> {
        let mut parts = [0u8, ..6];
        let mut splits = s.split(':');
        let mut i = 0;
        for split in splits {
            if i == 6 {
                return None;
            }
            match from_str_radix(split, 16) {
                Some(b) if split.len() != 0 => parts[i] = b,
                _ => return None,
            }
            i += 1;
        }

        if i == 6 {
            Some(MacAddr(parts[0], parts[1], parts[2], parts[3], parts[4], parts[5]))
        } else {
            None
        }
    }
}

#[test]
fn mac_addr_from_str() {
    assert_eq!("00:00:00:00:00:00".parse(), Some(MacAddr(0, 0, 0, 0, 0, 0)));
    assert_eq!("ff:ff:ff:ff:ff:ff".parse(), Some(MacAddr(0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF)));
    assert_eq!("12:34:56:78:90:ab".parse(), Some(MacAddr(0x12, 0x34, 0x56, 0x78, 0x90, 0xAB)));
    assert_eq!("::::::".parse::<MacAddr>(), None);
    assert_eq!("0::::::".parse::<MacAddr>(), None);
    assert_eq!("::::0::".parse::<MacAddr>(), None);
    assert_eq!("12:34:56:78".parse::<MacAddr>(), None);
    assert_eq!("12:34:56:78:".parse::<MacAddr>(), None);
    assert_eq!("12:34:56:78:90".parse::<MacAddr>(), None);
    assert_eq!("12:34:56:78:90:".parse::<MacAddr>(), None);
    assert_eq!("12:34:56:78:90:00:00".parse::<MacAddr>(), None);
    assert_eq!("xx:xx:xx:xx:xx:xx".parse::<MacAddr>(), None);
}

/// Represents a network interface and its associated addresses
#[deriving(Clone, PartialEq, Eq, Show)]
pub struct NetworkInterface {
    /// The name of the interface
    pub name: String,
    /// The interface index (operating system specific)
    pub index: u32,
    /// A MAC address for the interface
    pub mac: Option<MacAddr>,
    /// An IP addresses for the interface
    pub ips: Option<Vec<IpAddr>>,
    /// Operating system specific flags for the interface
    pub flags: u32,
}

impl NetworkInterface {
    /// Retreive the MAC address associated with the interface
    pub fn mac_address(&self) -> MacAddr {
        self.mac.unwrap()
    }

    /// Is the interface a loopback interface?
    pub fn is_loopback(&self) -> bool {
        self.flags & (libc::IFF_LOOPBACK as u32) != 0
    }
}

#[cfg(target_os = "linux")]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == libc::AF_PACKET {
            let sll: *const libc::sockaddr_ll = mem::transmute(sa);
            let mac = MacAddr((*sll).sll_addr[0], (*sll).sll_addr[1],
                              (*sll).sll_addr[2], (*sll).sll_addr[3],
                              (*sll).sll_addr[4], (*sll).sll_addr[5]);
            return (Some(mac), None);
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                        mem::size_of::<libc::sockaddr_storage>());
            return match addr {
                Ok(sa) => (None, Some(sa.ip)),
                Err(_) => (None, None)
            };
        }
    }
}

#[cfg(any(target_os = "freebsd", target_os = "macos"))]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use bindings::bpf;
    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == bpf::AF_LINK {
            let sdl: *const bpf::sockaddr_dl = mem::transmute(sa);
            let nlen = (*sdl).sdl_nlen as uint;
            let mac = MacAddr((*sdl).sdl_data[nlen + 0] as u8,
                              (*sdl).sdl_data[nlen + 1] as u8,
                              (*sdl).sdl_data[nlen + 2] as u8,
                              (*sdl).sdl_data[nlen + 3] as u8,
                              (*sdl).sdl_data[nlen + 4] as u8,
                              (*sdl).sdl_data[nlen + 5] as u8
                      );
            (Some(mac), None)
        } else {
            let addr = internal::sockaddr_to_addr(mem::transmute(sa),
                                        mem::size_of::<libc::sockaddr_storage>());
            match addr {
                Ok(sa) => (None, Some(sa.ip)),
                Err(_) => (None, None)
            }
        }
    }
}

/// Get a list of available network interfaces for the current machine.
#[inline]
pub fn get_network_interfaces() -> Vec<NetworkInterface> {
    get_network_interfaces_impl()
}

#[cfg(not(windows))]
fn get_network_interfaces_impl() -> Vec<NetworkInterface> {
    use std::string::String;

    let mut ifaces: Vec<NetworkInterface> = Vec::new();
    unsafe {
        let mut addrs: *mut libc::ifaddrs = mem::uninitialized();
        if libc::getifaddrs(&mut addrs) != 0 {
            return ifaces;
        }
        let mut addr = addrs;
        while addr.is_not_null() {
            let name = String::from_raw_buf((*addr).ifa_name as *const u8);
            let (mac, ip) = sockaddr_to_network_addr((*addr).ifa_addr as *const libc::sockaddr);
            let ni = NetworkInterface {
                name: name.clone(),
                index: 0,
                mac: mac,
                ips: ip.map(|ip| [ip].to_vec()),
                flags: (*addr).ifa_flags
            };
            let mut found: bool = false;
            for iface in ifaces.iter_mut() {
                if name == iface.name {
                    merge(iface, &ni);
                    found = true;
                }
            }
            if !found {
                ifaces.push(ni);
            }

            addr = (*addr).ifa_next;
        }
        libc::freeifaddrs(addrs);

        for iface in ifaces.iter_mut() {
            iface.index = iface.name.with_c_str(
                |name| libc::if_nametoindex(name)
            );
        }
        return ifaces;
    }

    fn merge(old: &mut NetworkInterface, new: &NetworkInterface) {
        old.mac = match new.mac {
            None => old.mac,
            _ => new.mac
        };
        match (&mut old.ips, &new.ips) {
            (&Some(ref mut old_ips), &Some(ref new_ips)) => old_ips.push_all(new_ips.as_slice()),
            _ => {}
        };
        old.flags = old.flags | new.flags;
    }

}

#[cfg(windows)]
fn get_network_interfaces_impl() -> Vec<NetworkInterface> {
    use std::str::from_utf8;
    use std::string::raw;

    use bindings::winpcap;

    let mut adapters_size = 0u32;

    unsafe {
        let mut tmp: winpcap::IP_ADAPTER_INFO = mem::zeroed();
        // FIXME [windows] This only gets IPv4 addresses - should use GetAdaptersAddresses
        winpcap::GetAdaptersInfo(
            &mut tmp,
            &mut adapters_size
        );
    }

    let vec_size = adapters_size / mem::size_of::<winpcap::IP_ADAPTER_INFO>() as u32;

    let mut adapters = Vec::with_capacity(vec_size as uint);

    // FIXME [windows] Check return code
    unsafe {
        winpcap::GetAdaptersInfo(adapters.as_mut_ptr(), &mut adapters_size);
    }

    // Create a complete list of NetworkInterfaces for the machine
    let mut cursor = adapters.as_mut_ptr();
    let mut all_ifaces = Vec::with_capacity(vec_size as uint);
    while cursor.is_not_null() {
        let mac = unsafe {
                    MacAddr((*cursor).Address[0],
                            (*cursor).Address[1],
                            (*cursor).Address[2],
                            (*cursor).Address[3],
                            (*cursor).Address[4],
                            (*cursor).Address[5])
                  };
        let mut ip_cursor = unsafe { &mut (*cursor).IpAddressList as winpcap::PIP_ADDR_STRING};
        let mut ips: Vec<IpAddr> = Vec::new();
        while ip_cursor.is_not_null() {
            let ip_str = unsafe {
                            raw::from_buf((*ip_cursor).IpAddress.String.as_ptr() as *const u8)
                         };
            ips.push(from_str(ip_str.as_slice()).unwrap());
            ip_cursor = unsafe { (*ip_cursor).Next };
        }
        unsafe {
            all_ifaces.push(NetworkInterface {
                        name: raw::from_buf((*cursor).AdapterName.as_ptr() as *const u8),
                        index: (*cursor).Index,
                        mac: Some(mac),
                        ips: Some(ips),
                        //flags: (*cursor).Type, // FIXME [windows]
                        flags: 0,
                     });

            cursor = (*cursor).Next;
        }
    }

    let mut buf = [0u8, ..4096];
    let mut buflen = buf.len() as u32;

    // Gets list of supported adapters in form:
    // adapter1\0adapter2\0\0desc1\0desc2\0\0
    if unsafe { winpcap::PacketGetAdapterNames(buf.as_mut_ptr() as *mut i8, &mut buflen) } == 0 {
        // FIXME [windows] Should allocate a buffer big enough and try again
        //        - size should be buf.len() + buflen (buflen is overwritten)
        panic!("FIXME [windows] unable to get interface list");
    }

    let buf_str = from_utf8(buf).unwrap();
    let iface_names = buf_str.split_str("\0\0").next();
    let mut vec = Vec::new();

    // Return only supported adapters
    match iface_names {
        Some(iface_names) => {
            for iface in iface_names.split('\0') {
                let name = iface.to_string();
                let next = all_ifaces.iter()
                                     .filter(|x| name.as_slice().ends_with(x.name.as_slice()))
                                     .next();
                if next.is_some() {
                    let mut iface = next.unwrap().clone();
                    iface.name = name;
                    vec.push(iface);
                }
            }
        },
        None => ()
    };

    vec
}

<|endoftext|>


<commit_before>//! Module: czmq-zmsg

use {czmq_sys, Error, ErrorKind, RawInterface, Result, Sockish, ZFrame};
use std::{error, fmt, mem, ptr, result};
use std::ffi::{CStr, CString};

pub struct ZMsg {
    zmsg: *mut czmq_sys::zmsg_t,
    owned: bool,
}

unsafe impl Send for ZMsg {}
unsafe impl Sync for ZMsg {}

impl Drop for ZMsg {
    fn drop(&mut self) {
        if self.owned {
            unsafe { czmq_sys::zmsg_destroy(&mut self.zmsg) };
        }
    }
}

impl ZMsg {
    pub fn new() -> ZMsg {
        ZMsg {
            zmsg: unsafe { czmq_sys::zmsg_new() },
            owned: true,
        }
    }

    pub fn recv<S: Sockish>(source: &S) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_recv(source.borrow_raw()) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    // XXX We'll have to roll our own here as we can't imitate a C
    // file handle without C boilerplate...which we're not doing!
    // pub fn zmsg_load(file: *mut FILE) -> *mut zmsg_t;

    pub fn encode(&self) -> Result<ZFrame> {
        let zframe = unsafe { czmq_sys::zmsg_encode(self.zmsg) };

        if zframe == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZFrame::from_raw(zframe, true))
        }
    }

    pub fn decode(frame: &ZFrame) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_decode(frame.borrow_raw()) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    pub fn new_signal(status: u8) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_new_signal(status) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    pub fn send<D: Sockish>(self, dest: &D) -> Result<()> {
        let mut zmsg = self;

        let rc = unsafe { czmq_sys::zmsg_send(&mut zmsg.zmsg, dest.borrow_raw()) };
        if rc == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(())
        }
    }

    // pub fn zmsg_sendm(self_p: *mut *mut zmsg_t,
    //                   dest: *mut ::std::os::raw::c_void)
    //  -> ::std::os::raw::c_int;

    pub fn size(&self) -> usize {
        unsafe { czmq_sys::zmsg_size(self.zmsg) as usize }
    }

    // pub fn zmsg_content_size(_self: *mut zmsg_t) -> size_t;
    // pub fn zmsg_prepend(_self: *mut zmsg_t, frame_p: *mut *mut zframe_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_append(_self: *mut zmsg_t, frame_p: *mut *mut zframe_t)
    //  -> ::std::os::raw::c_int;

    pub fn pop(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_pop(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, true))
        }
    }

    // pub fn zmsg_pushmem(_self: *mut zmsg_t,
    //                     src: *const ::std::os::raw::c_void, size: size_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_addmem(_self: *mut zmsg_t, src: *const ::std::os::raw::c_void,
    //                    size: size_t) -> ::std::os::raw::c_int;
    // pub fn zmsg_pushstr(_self: *mut zmsg_t,
    //                     string: *const ::std::os::raw::c_char)
    //  -> ::std::os::raw::c_int;

    pub fn addstr(&self, string: &str) -> Result<()> {
        let string_c = CString::new(string).unwrap_or(CString::new("").unwrap());
        let rc = unsafe { czmq_sys::zmsg_addstr(self.zmsg, string_c.as_ptr()) };

        // Deliberately leak this memory, which will be managed by C
        mem::forget(string_c);

        if rc == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(())
        }
    }

    // pub fn zmsg_pushstrf(_self: *mut zmsg_t,
    //                      format: *const ::std::os::raw::c_char, ...)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_addstrf(_self: *mut zmsg_t,
    //                     format: *const ::std::os::raw::c_char, ...)
    //  -> ::std::os::raw::c_int;

    pub fn popstr(&self) -> Option<result::Result<String, Vec<u8>>> {
        let ptr = unsafe { czmq_sys::zmsg_popstr(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            let bytes = c_string.as_bytes().to_vec();
            match c_string.into_string() {
                Ok(s) => Some(Ok(s)),
                Err(_) => Some(Err(bytes)),
            }
        }
    }

    pub fn popbytes(&self) -> Option<Vec<u8>> {
        let ptr = unsafe { czmq_sys::zmsg_popstr(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            Some(c_string.to_bytes().to_vec())
        }
    }

    // pub fn zmsg_addmsg(_self: *mut zmsg_t, msg_p: *mut *mut zmsg_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_popmsg(_self: *mut zmsg_t) -> *mut zmsg_t;
    // pub fn zmsg_remove(_self: *mut zmsg_t, frame: *mut zframe_t);

    pub fn first(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_first(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    pub fn next(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_next(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    // We can't call this fn last() as it conflicts with
    // Iterator::last(), which is implemented for ZMsg.
    pub fn ref_last(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_last(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    // pub fn zmsg_save(_self: *mut zmsg_t, file: *mut FILE)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_dup(_self: *mut zmsg_t) -> *mut zmsg_t;
    // pub fn zmsg_print(_self: *mut zmsg_t);
    // pub fn zmsg_eq(_self: *mut zmsg_t, other: *mut zmsg_t) -> u8;

    pub fn signal(&self) -> Result<u8> {
        let signal = unsafe { czmq_sys::zmsg_signal(self.zmsg) };
        if signal == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(signal as u8)
        }
    }

    // pub fn zmsg_test(verbose: u8);

    pub fn borrow_raw(&self) -> *mut czmq_sys::zmsg_t {
        self.zmsg
    }
}

impl RawInterface<czmq_sys::zmsg_t> for ZMsg {
    fn from_raw(ptr: *mut czmq_sys::zmsg_t, owned: bool) -> ZMsg {
        ZMsg {
            zmsg: ptr,
            owned: owned,
        }
    }

    fn into_raw(mut self) -> *mut czmq_sys::zmsg_t {
        self.owned = false;
        self.zmsg
    }

    fn borrow_raw(&self) -> *mut czmq_sys::zmsg_t {
        self.zmsg
    }
}

impl Iterator for ZMsg {
    type Item = ZFrame;

    fn next(&mut self) -> Option<Self::Item> {
        ZMsg::next(self)
    }
}

#[derive(Debug)]
pub enum ZMsgError {
    CmdFailed,
}

impl fmt::Display for ZMsgError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ZMsgError::CmdFailed => write!(f, "ZMsg command failed"),
        }
    }
}

impl error::Error for ZMsgError {
    fn description(&self) -> &str {
        match *self {
            ZMsgError::CmdFailed => "ZMsg command failed",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use {zmq, ZSock, zsys_init};

    #[test]
    fn test_sendrecv_zmq() {
        let mut ctx = zmq::Context::new();

        let mut server = ctx.socket(zmq::REP).unwrap();
        server.bind("inproc://zmsg_sendrecv_zmq").unwrap();

        let mut client = ctx.socket(zmq::REQ).unwrap();
        client.connect("inproc://zmsg_sendrecv_zmq").unwrap();

        let zmsg = ZMsg::new();
        zmsg.addstr("Hello world!").unwrap();
        zmsg.send(&client).unwrap();

        let zmsg_recv = ZMsg::recv(&server).unwrap();
        assert_eq!(zmsg_recv.popstr().unwrap().unwrap(), "Hello world!");
    }

    #[test]
    fn test_sendrecv_zsock() {
        zsys_init();

        let server = ZSock::new_rep("inproc://zmsg_sendrecv_zsock").unwrap();
        let client = ZSock::new_req("inproc://zmsg_sendrecv_zsock").unwrap();

        let zmsg = ZMsg::new();
        zmsg.addstr("Hello world!").unwrap();
        zmsg.send(&client).unwrap();

        let msg = ZMsg::recv(&server).unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "Hello world!");
    }

    #[test]
    fn test_encode_decode() {
        let msg = ZMsg::new();
        msg.addstr("moo").unwrap();
        let msg_decoded = ZMsg::decode(&msg.encode().unwrap()).unwrap();
        assert_eq!(msg_decoded.popstr().unwrap().unwrap(), "moo");
    }

    #[test]
    fn test_signal() {
        let msg = ZMsg::new_signal(97).unwrap();
        assert_eq!(msg.signal().unwrap(), 97);
    }

    #[test]
    fn test_size() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.size(), 1);
    }

    #[test]
    fn test_pop() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.pop().unwrap().data().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_popstr() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_popbytes() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.popbytes().unwrap(), "123".as_bytes());
    }

    #[test]
    fn test_iter_fns() {
        let msg = ZMsg::new();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.addstr("3").unwrap();

        let frame = msg.first().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "1");

        let frame = msg.next().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "2");

        let frame = msg.ref_last().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "3");
    }

    #[test]
    fn test_iter() {
        let msg = ZMsg::new();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.addstr("3").unwrap();

        let mut i = 1;
        for x in msg {
            assert_eq!(x.data().unwrap().unwrap(), i.to_string());
            i += 1;
        }
    }
}
<commit_msg>Implement prepend/append fns for ZMsg<commit_after>//! Module: czmq-zmsg

use {czmq_sys, Error, ErrorKind, RawInterface, Result, Sockish, ZFrame};
use std::{error, fmt, mem, ptr, result};
use std::ffi::{CStr, CString};

pub struct ZMsg {
    zmsg: *mut czmq_sys::zmsg_t,
    owned: bool,
}

unsafe impl Send for ZMsg {}
unsafe impl Sync for ZMsg {}

impl Drop for ZMsg {
    fn drop(&mut self) {
        if self.owned {
            unsafe { czmq_sys::zmsg_destroy(&mut self.zmsg) };
        }
    }
}

impl ZMsg {
    pub fn new() -> ZMsg {
        ZMsg {
            zmsg: unsafe { czmq_sys::zmsg_new() },
            owned: true,
        }
    }

    pub fn recv<S: Sockish>(source: &S) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_recv(source.borrow_raw()) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    // XXX We'll have to roll our own here as we can't imitate a C
    // file handle without C boilerplate...which we're not doing!
    // pub fn zmsg_load(file: *mut FILE) -> *mut zmsg_t;

    pub fn encode(&self) -> Result<ZFrame> {
        let zframe = unsafe { czmq_sys::zmsg_encode(self.zmsg) };

        if zframe == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZFrame::from_raw(zframe, true))
        }
    }

    pub fn decode(frame: &ZFrame) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_decode(frame.borrow_raw()) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    pub fn new_signal(status: u8) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_new_signal(status) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    pub fn send<D: Sockish>(self, dest: &D) -> Result<()> {
        let mut zmsg = self;

        let rc = unsafe { czmq_sys::zmsg_send(&mut zmsg.zmsg, dest.borrow_raw()) };
        if rc == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(())
        }
    }

    // pub fn zmsg_sendm(self_p: *mut *mut zmsg_t,
    //                   dest: *mut ::std::os::raw::c_void)
    //  -> ::std::os::raw::c_int;

    pub fn size(&self) -> usize {
        unsafe { czmq_sys::zmsg_size(self.zmsg) as usize }
    }

    // pub fn zmsg_content_size(_self: *mut zmsg_t) -> size_t;

    pub fn prepend(&self, frame: ZFrame) -> Result<()> {
        let rc = unsafe { czmq_sys::zmsg_prepend(self.zmsg, &mut frame.into_raw()) };

        if rc == 0 {
            Ok(())
        } else {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        }
    }

    pub fn append(&self, frame: ZFrame) -> Result<()> {
        let rc = unsafe { czmq_sys::zmsg_append(self.zmsg, &mut frame.into_raw()) };

        if rc == 0 {
            Ok(())
        } else {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        }
    }

    pub fn pop(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_pop(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, true))
        }
    }

    // pub fn zmsg_pushmem(_self: *mut zmsg_t,
    //                     src: *const ::std::os::raw::c_void, size: size_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_addmem(_self: *mut zmsg_t, src: *const ::std::os::raw::c_void,
    //                    size: size_t) -> ::std::os::raw::c_int;
    // pub fn zmsg_pushstr(_self: *mut zmsg_t,
    //                     string: *const ::std::os::raw::c_char)
    //  -> ::std::os::raw::c_int;

    pub fn addstr(&self, string: &str) -> Result<()> {
        let string_c = CString::new(string).unwrap_or(CString::new("").unwrap());
        let rc = unsafe { czmq_sys::zmsg_addstr(self.zmsg, string_c.as_ptr()) };

        // Deliberately leak this memory, which will be managed by C
        mem::forget(string_c);

        if rc == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(())
        }
    }

    // pub fn zmsg_pushstrf(_self: *mut zmsg_t,
    //                      format: *const ::std::os::raw::c_char, ...)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_addstrf(_self: *mut zmsg_t,
    //                     format: *const ::std::os::raw::c_char, ...)
    //  -> ::std::os::raw::c_int;

    pub fn popstr(&self) -> Option<result::Result<String, Vec<u8>>> {
        let ptr = unsafe { czmq_sys::zmsg_popstr(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            let bytes = c_string.as_bytes().to_vec();
            match c_string.into_string() {
                Ok(s) => Some(Ok(s)),
                Err(_) => Some(Err(bytes)),
            }
        }
    }

    pub fn popbytes(&self) -> Option<Vec<u8>> {
        let ptr = unsafe { czmq_sys::zmsg_popstr(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            Some(c_string.to_bytes().to_vec())
        }
    }

    // pub fn zmsg_addmsg(_self: *mut zmsg_t, msg_p: *mut *mut zmsg_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_popmsg(_self: *mut zmsg_t) -> *mut zmsg_t;
    // pub fn zmsg_remove(_self: *mut zmsg_t, frame: *mut zframe_t);

    pub fn first(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_first(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    pub fn next(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_next(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    // We can't call this fn last() as it conflicts with
    // Iterator::last(), which is implemented for ZMsg.
    pub fn ref_last(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_last(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    // pub fn zmsg_save(_self: *mut zmsg_t, file: *mut FILE)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_dup(_self: *mut zmsg_t) -> *mut zmsg_t;
    // pub fn zmsg_print(_self: *mut zmsg_t);
    // pub fn zmsg_eq(_self: *mut zmsg_t, other: *mut zmsg_t) -> u8;

    pub fn signal(&self) -> Result<u8> {
        let signal = unsafe { czmq_sys::zmsg_signal(self.zmsg) };
        if signal == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(signal as u8)
        }
    }

    // pub fn zmsg_test(verbose: u8);

    pub fn borrow_raw(&self) -> *mut czmq_sys::zmsg_t {
        self.zmsg
    }
}

impl RawInterface<czmq_sys::zmsg_t> for ZMsg {
    fn from_raw(ptr: *mut czmq_sys::zmsg_t, owned: bool) -> ZMsg {
        ZMsg {
            zmsg: ptr,
            owned: owned,
        }
    }

    fn into_raw(mut self) -> *mut czmq_sys::zmsg_t {
        self.owned = false;
        self.zmsg
    }

    fn borrow_raw(&self) -> *mut czmq_sys::zmsg_t {
        self.zmsg
    }
}

impl Iterator for ZMsg {
    type Item = ZFrame;

    fn next(&mut self) -> Option<Self::Item> {
        ZMsg::next(self)
    }
}

#[derive(Debug)]
pub enum ZMsgError {
    CmdFailed,
}

impl fmt::Display for ZMsgError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ZMsgError::CmdFailed => write!(f, "ZMsg command failed"),
        }
    }
}

impl error::Error for ZMsgError {
    fn description(&self) -> &str {
        match *self {
            ZMsgError::CmdFailed => "ZMsg command failed",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use {zmq, ZFrame, ZSock, zsys_init};

    #[test]
    fn test_sendrecv_zmq() {
        let mut ctx = zmq::Context::new();

        let mut server = ctx.socket(zmq::REP).unwrap();
        server.bind("inproc://zmsg_sendrecv_zmq").unwrap();

        let mut client = ctx.socket(zmq::REQ).unwrap();
        client.connect("inproc://zmsg_sendrecv_zmq").unwrap();

        let zmsg = ZMsg::new();
        zmsg.addstr("Hello world!").unwrap();
        zmsg.send(&client).unwrap();

        let zmsg_recv = ZMsg::recv(&server).unwrap();
        assert_eq!(zmsg_recv.popstr().unwrap().unwrap(), "Hello world!");
    }

    #[test]
    fn test_sendrecv_zsock() {
        zsys_init();

        let server = ZSock::new_rep("inproc://zmsg_sendrecv_zsock").unwrap();
        let client = ZSock::new_req("inproc://zmsg_sendrecv_zsock").unwrap();

        let zmsg = ZMsg::new();
        zmsg.addstr("Hello world!").unwrap();
        zmsg.send(&client).unwrap();

        let msg = ZMsg::recv(&server).unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "Hello world!");
    }

    #[test]
    fn test_encode_decode() {
        let msg = ZMsg::new();
        msg.addstr("moo").unwrap();
        let msg_decoded = ZMsg::decode(&msg.encode().unwrap()).unwrap();
        assert_eq!(msg_decoded.popstr().unwrap().unwrap(), "moo");
    }

    #[test]
    fn test_signal() {
        let msg = ZMsg::new_signal(97).unwrap();
        assert_eq!(msg.signal().unwrap(), 97);
    }

    #[test]
    fn test_size() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.size(), 1);
    }

    #[test]
    fn test_prepend() {
        let msg = ZMsg::new();
        msg.prepend(ZFrame::from("123").unwrap()).unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_append() {
        let msg = ZMsg::new();
        msg.append(ZFrame::from("123").unwrap()).unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_pop() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.pop().unwrap().data().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_popstr() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_popbytes() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.popbytes().unwrap(), "123".as_bytes());
    }

    #[test]
    fn test_iter_fns() {
        let msg = ZMsg::new();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.addstr("3").unwrap();

        let frame = msg.first().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "1");

        let frame = msg.next().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "2");

        let frame = msg.ref_last().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "3");
    }

    #[test]
    fn test_iter() {
        let msg = ZMsg::new();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.addstr("3").unwrap();

        let mut i = 1;
        for x in msg {
            assert_eq!(x.data().unwrap().unwrap(), i.to_string());
            i += 1;
        }
    }
}
<|endoftext|>


<commit_before>use std::io::net::ip::{Ipv4Addr, SocketAddr};

mod libtorresmo {
    extern crate time;

    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::SocketAddr;
    use std::io::IoResult;
    use std::mem::{to_be32, to_be16, transmute};
    use std::rand::random;

    #[allow(dead_code,non_camel_case_types)]
    enum UtpPacketType {
        ST_DATA  = 0,
        ST_FIN   = 1,
        ST_STATE = 2,
        ST_RESET = 3,
        ST_SYN   = 4,
    }

    #[allow(dead_code)]
    #[packed]
    struct UtpPacketHeader {
        type_ver: u8, // type: u4, ver: u4
        extension: u8,
        connection_id: u16,
        timestamp_microseconds: u32,
        timestamp_difference_microseconds: u32,
        wnd_size: u32,
        seq_nr: u16,
        ack_nr: u16,
    }

    impl UtpPacketHeader {
        fn bytes(&self) -> &[u8] {
            unsafe {
                let buf: &[u8, ..20] = transmute(self);
                return buf.as_slice();
            }
        }

        fn len(&self) -> uint {
            static HEADER_SIZE: uint = 20;
            return HEADER_SIZE;
        }
    }

    #[allow(dead_code)]
    struct UtpPacket {
        header: UtpPacketHeader,
        payload: Vec<u8>,
    }

    impl UtpPacket {
        /// Constructs a new, empty UtpPacket.
        fn new() -> UtpPacket {
            UtpPacket {
                header: UtpPacketHeader {
                    type_ver: ST_DATA as u8 << 4 | 1,
                    extension: 0,
                    connection_id: 0,
                    timestamp_microseconds: 0,
                    timestamp_difference_microseconds: 0,
                    wnd_size: 0,
                    seq_nr: 0,
                    ack_nr: 0,
                },
                payload: Vec::new(),
            }
        }

        fn set_type(&mut self, t: UtpPacketType) {
            let version = 0x0F & self.header.type_ver;
            self.header.type_ver = t as u8 << 4 | version;
        }

        fn bytes(&self) -> Vec<u8> {
            let mut buf: Vec<u8> = Vec::with_capacity(self.len());
            buf.push_all(self.header.bytes());
            buf.push_all(self.payload.as_slice());
            return buf;
        }

        fn len(&self) -> uint {
            self.header.len() + self.payload.len()
        }
    }

    pub struct UtpSocket {
        socket: UdpSocket,
    }

    impl UtpSocket {
        pub fn bind(addr: SocketAddr) -> IoResult<UtpSocket> {
            let skt = UdpSocket::bind(addr);
            match skt {
                Ok(x)  => Ok(UtpSocket { socket: x }),
                Err(e) => Err(e)
            }
        }

        pub fn connect(self, other: SocketAddr) -> UtpStream {
            let stream = UtpStream::new(self.socket, other);
            stream.connect()
        }

        pub fn recvfrom(&mut self, buf: &mut[u8]) -> IoResult<(uint,SocketAddr)> {
            self.socket.recvfrom(buf)
        }
    }

    impl Clone for UtpSocket {
        fn clone(&self) -> UtpSocket {
            UtpSocket {
                socket: self.socket.clone(),
            }
        }
    }

    #[allow(dead_code)]
    pub struct UtpStream {
        socket: UdpSocket,
        connected_to: SocketAddr,
        connection_id: u16,
        seq_nr: u16,
        ack_nr: u16,
    }

    #[allow(dead_code)]
    impl UtpStream {
        pub fn new(socket: UdpSocket, conn: SocketAddr) -> UtpStream {
            UtpStream {
                socket: socket,
                connected_to: conn,
                connection_id: to_be16(random()),
                seq_nr: 1,
                ack_nr: 0,
            }
        }

        pub fn connect(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = to_be16(self.seq_nr);
            let t = time::get_time();
            packet.header.timestamp_microseconds = to_be32((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32);

            // Send packet
            let _result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;

            self
        }

        pub fn disconnect(self) -> UdpSocket {
            self.socket
        }
    }

    impl Reader for UtpStream {
        fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
            match self.socket.recvfrom(buf) {
                Ok((_nread, src)) if src != self.connected_to => Ok(0),
                Ok((nread, _src)) => Ok(nread),
                Err(e) => Err(e),
            }
        }
    }
    impl Writer for UtpStream {
        fn write(&mut self, buf: &[u8]) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = to_be16(self.seq_nr);
            let t = time::get_time();
            packet.header.timestamp_microseconds = to_be32((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32);

            // Send packet
            let result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;
            result
        }
    }
}

fn usage() {
    println!("Usage: torresmo [-s|-c] <address> <port>");
}

fn main() {
    use libtorresmo::{UtpSocket, UtpStream};
    use std::from_str::FromStr;

    // Defaults
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

    if args.len() == 4 {
        let ip = FromStr::from_str(args.get(2).as_slice()).unwrap();
        let port = FromStr::from_str(args.get(3).as_slice()).unwrap();
        addr = SocketAddr { ip: ip, port: port };
    }

    match args.get(1).as_slice() {
        "-s" => {
            let mut buf = [0, ..512];
            let sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                let mut sock = sock.clone();
                match sock.recvfrom(buf) {
                    Ok((_, src)) => {
                        spawn(proc() {
                            let mut stream = sock.connect(src);
                            let payload = String::from_str("Hello\n").into_bytes();

                            // Send uTP packet
                            let _ = stream.write(payload.as_slice());
                        })
                    }
                    Err(_) => {}
                }
            }
        }
        "-c" => {
            let sock = UtpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: std::rand::random() }).unwrap();
            let mut stream = sock.connect(addr);
            let _ = stream.write([0]);
            drop(stream);
        }
        _ => usage(),
    }
}
<commit_msg>Add string representation for UtpPacketHeader.<commit_after>use std::io::net::ip::{Ipv4Addr, SocketAddr};

mod libtorresmo {
    extern crate time;

    use std::io::net::udp::UdpSocket;
    use std::io::net::ip::SocketAddr;
    use std::io::IoResult;
    use std::mem::{to_be32, to_be16, transmute};
    use std::rand::random;
    use std::fmt;

    #[allow(dead_code,non_camel_case_types)]
    enum UtpPacketType {
        ST_DATA  = 0,
        ST_FIN   = 1,
        ST_STATE = 2,
        ST_RESET = 3,
        ST_SYN   = 4,
    }

    #[allow(dead_code)]
    #[packed]
    struct UtpPacketHeader {
        type_ver: u8, // type: u4, ver: u4
        extension: u8,
        connection_id: u16,
        timestamp_microseconds: u32,
        timestamp_difference_microseconds: u32,
        wnd_size: u32,
        seq_nr: u16,
        ack_nr: u16,
    }

    impl UtpPacketHeader {
        fn bytes(&self) -> &[u8] {
            unsafe {
                let buf: &[u8, ..20] = transmute(self);
                return buf.as_slice();
            }
        }

        fn len(&self) -> uint {
            static HEADER_SIZE: uint = 20;
            return HEADER_SIZE;
        }
    }

    impl fmt::Show for UtpPacketHeader {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            let t: UtpPacketType = unsafe { transmute(self.type_ver >> 4) };
            let ptype = match t {
                ST_DATA  => "ST_DATA",
                ST_FIN   => "ST_FIN",
                ST_STATE => "ST_STATE",
                ST_RESET => "ST_RESET",
                ST_SYN   => "ST_SYN",
            };
            write!(f, "(type: {}, version: {}, extension: {}, connection_id: {}, timestamp_microseconds: {}, timestamp_difference_microseconds: {}, wnd_size: {}, seq_nr: {}, ack_nr: {})", ptype, self.type_ver & 0x0F, self.extension, self.connection_id, self.timestamp_microseconds, self.timestamp_difference_microseconds, self.wnd_size, self.seq_nr, self.ack_nr)
        }
    }

    #[allow(dead_code)]
    struct UtpPacket {
        header: UtpPacketHeader,
        payload: Vec<u8>,
    }

    impl UtpPacket {
        /// Constructs a new, empty UtpPacket.
        fn new() -> UtpPacket {
            UtpPacket {
                header: UtpPacketHeader {
                    type_ver: ST_DATA as u8 << 4 | 1,
                    extension: 0,
                    connection_id: 0,
                    timestamp_microseconds: 0,
                    timestamp_difference_microseconds: 0,
                    wnd_size: 0,
                    seq_nr: 0,
                    ack_nr: 0,
                },
                payload: Vec::new(),
            }
        }

        fn set_type(&mut self, t: UtpPacketType) {
            let version = 0x0F & self.header.type_ver;
            self.header.type_ver = t as u8 << 4 | version;
        }

        fn bytes(&self) -> Vec<u8> {
            let mut buf: Vec<u8> = Vec::with_capacity(self.len());
            buf.push_all(self.header.bytes());
            buf.push_all(self.payload.as_slice());
            return buf;
        }

        fn len(&self) -> uint {
            self.header.len() + self.payload.len()
        }
    }

    pub struct UtpSocket {
        socket: UdpSocket,
    }

    impl UtpSocket {
        pub fn bind(addr: SocketAddr) -> IoResult<UtpSocket> {
            let skt = UdpSocket::bind(addr);
            match skt {
                Ok(x)  => Ok(UtpSocket { socket: x }),
                Err(e) => Err(e)
            }
        }

        pub fn connect(self, other: SocketAddr) -> UtpStream {
            let stream = UtpStream::new(self.socket, other);
            stream.connect()
        }

        pub fn recvfrom(&mut self, buf: &mut[u8]) -> IoResult<(uint,SocketAddr)> {
            self.socket.recvfrom(buf)
        }
    }

    impl Clone for UtpSocket {
        fn clone(&self) -> UtpSocket {
            UtpSocket {
                socket: self.socket.clone(),
            }
        }
    }

    #[allow(dead_code)]
    pub struct UtpStream {
        socket: UdpSocket,
        connected_to: SocketAddr,
        connection_id: u16,
        seq_nr: u16,
        ack_nr: u16,
    }

    #[allow(dead_code)]
    impl UtpStream {
        pub fn new(socket: UdpSocket, conn: SocketAddr) -> UtpStream {
            UtpStream {
                socket: socket,
                connected_to: conn,
                connection_id: to_be16(random()),
                seq_nr: 1,
                ack_nr: 0,
            }
        }

        pub fn connect(mut self) -> UtpStream {
            let mut packet = UtpPacket::new();
            packet.set_type(ST_SYN);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = to_be16(self.seq_nr);
            let t = time::get_time();
            packet.header.timestamp_microseconds = to_be32((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32);

            // Send packet
            let _result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;

            self
        }

        pub fn disconnect(self) -> UdpSocket {
            self.socket
        }
    }

    impl Reader for UtpStream {
        fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
            match self.socket.recvfrom(buf) {
                Ok((_nread, src)) if src != self.connected_to => Ok(0),
                Ok((nread, _src)) => Ok(nread),
                Err(e) => Err(e),
            }
        }
    }
    impl Writer for UtpStream {
        fn write(&mut self, buf: &[u8]) -> IoResult<()> {
            let mut packet = UtpPacket::new();
            packet.payload = Vec::from_slice(buf);
            packet.header.connection_id = self.connection_id;
            packet.header.seq_nr = to_be16(self.seq_nr);
            let t = time::get_time();
            packet.header.timestamp_microseconds = to_be32((t.sec * 1_000_000) as u32 + (t.nsec/1000) as u32);

            // Send packet
            let result = self.socket.sendto(packet.bytes().as_slice(), self.connected_to);

            self.seq_nr += 1;
            result
        }
    }
}

fn usage() {
    println!("Usage: torresmo [-s|-c] <address> <port>");
}

fn main() {
    use libtorresmo::{UtpSocket, UtpStream};
    use std::from_str::FromStr;

    // Defaults
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

    if args.len() == 4 {
        let ip = FromStr::from_str(args.get(2).as_slice()).unwrap();
        let port = FromStr::from_str(args.get(3).as_slice()).unwrap();
        addr = SocketAddr { ip: ip, port: port };
    }

    match args.get(1).as_slice() {
        "-s" => {
            let mut buf = [0, ..512];
            let sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                let mut sock = sock.clone();
                match sock.recvfrom(buf) {
                    Ok((_, src)) => {
                        spawn(proc() {
                            let mut stream = sock.connect(src);
                            let payload = String::from_str("Hello\n").into_bytes();

                            // Send uTP packet
                            let _ = stream.write(payload.as_slice());
                        })
                    }
                    Err(_) => {}
                }
            }
        }
        "-c" => {
            let sock = UtpSocket::bind(SocketAddr { ip: Ipv4Addr(127,0,0,1), port: std::rand::random() }).unwrap();
            let mut stream = sock.connect(addr);
            let _ = stream.write([0]);
            drop(stream);
        }
        _ => usage(),
    }
}
<|endoftext|>


<commit_before>//! CLI tool to read Cranelift IR files and compile them into native code.

use cranelift_codegen::isa::TargetIsa;
use cranelift_codegen::print_errors::pretty_error;
use cranelift_codegen::settings::FlagsOrIsa;
use cranelift_codegen::Context;
use cranelift_codegen::{binemit, ir};
use cranelift_reader::parse_test;
use std::path::Path;
use std::path::PathBuf;
use utils::{parse_sets_and_triple, read_to_string};

struct PrintRelocs {
    flag_print: bool,
}

impl binemit::RelocSink for PrintRelocs {
    fn reloc_ebb(
        &mut self,
        where_: binemit::CodeOffset,
        r: binemit::Reloc,
        offset: binemit::CodeOffset,
    ) {
        if self.flag_print {
            println!("reloc_ebb: {} {} at {}", r, offset, where_);
        }
    }

    fn reloc_external(
        &mut self,
        where_: binemit::CodeOffset,
        r: binemit::Reloc,
        name: &ir::ExternalName,
        addend: binemit::Addend,
    ) {
        if self.flag_print {
            println!("reloc_ebb: {} {} {} at {}", r, name, addend, where_);
        }
    }

    fn reloc_jt(&mut self, where_: binemit::CodeOffset, r: binemit::Reloc, jt: ir::JumpTable) {
        if self.flag_print {
            println!("reloc_ebb: {} {} at {}", r, jt, where_);
        }
    }
}

struct PrintTraps {
    flag_print: bool,
}

impl binemit::TrapSink for PrintTraps {
    fn trap(&mut self, offset: binemit::CodeOffset, _srcloc: ir::SourceLoc, code: ir::TrapCode) {
        if self.flag_print {
            println!("trap: {} at {}", code, offset);
        }
    }
}

pub fn run(
    files: Vec<String>,
    flag_print: bool,
    flag_set: &[String],
    flag_isa: &str,
) -> Result<(), String> {
    let parsed = parse_sets_and_triple(flag_set, flag_isa)?;

    for filename in files {
        let path = Path::new(&filename);
        let name = String::from(path.as_os_str().to_string_lossy());
        handle_module(flag_print, &path.to_path_buf(), &name, parsed.as_fisa())?;
    }
    Ok(())
}

fn handle_module(
    flag_print: bool,
    path: &PathBuf,
    name: &str,
    fisa: FlagsOrIsa,
) -> Result<(), String> {
    let buffer = read_to_string(&path).map_err(|e| format!("{}: {}", name, e))?;
    let test_file = parse_test(&buffer, None, None).map_err(|e| format!("{}: {}", name, e))?;

    // If we have an isa from the command-line, use that. Otherwise if the
    // file contains a unique isa, use that.
    let isa = if let Some(isa) = fisa.isa {
        isa
    } else if let Some(isa) = test_file.isa_spec.unique_isa() {
        isa
    } else {
        return Err(String::from("compilation requires a target isa"));
    };

    for (func, _) in test_file.functions {
        let mut context = Context::new();
        context.func = func;

        // Compile and encode the result to machine code.
        let total_size = context
            .compile(isa)
            .map_err(|err| pretty_error(&context.func, Some(isa), err))?;

        let mut mem = Vec::new();
        mem.resize(total_size as usize, 0);

        let mut relocs = PrintRelocs { flag_print };
        let mut traps = PrintTraps { flag_print };
        let mut code_sink: binemit::MemoryCodeSink;
        unsafe {
            code_sink = binemit::MemoryCodeSink::new(mem.as_mut_ptr(), &mut relocs, &mut traps);
        }
        isa.emit_function_to_memory(&context.func, &mut code_sink);

        if flag_print {
            println!("{}", context.func.display(isa));
        }

        if flag_print {
            print!(".byte ");
            let mut first = true;
            for byte in &mem {
                if first {
                    first = false;
                } else {
                    print!(", ");
                }
                print!("{}", byte);
            }

            println!();
            print_disassembly(isa, &mem[0..code_sink.code_size as usize])?;
            print_readonly_data(&mem[code_sink.code_size as usize..total_size as usize]);
        }
    }

    Ok(())
}

fn print_readonly_data(mem: &[u8]) {
    if mem.len() == 0 {
        return;
    }

    println!("\nFollowed by {} bytes of read-only data:", mem.len());

    for (i, byte) in mem.iter().enumerate() {
        if i % 16 == 0 {
            if i != 0 {
                println!();
            }
            print!("{:4}: ", i);
        }
        if i % 4 == 0 {
            print!(" ");
        }
        print!("{:02x} ", byte);
    }
    println!();
}

cfg_if! {
    if #[cfg(feature = "disas")] {
        use capstone::prelude::*;
        use target_lexicon::Architecture;
        use std::fmt::Write;

        fn get_disassembler(isa: &TargetIsa) -> Result<Capstone, String> {
            let cs = match isa.triple().architecture {
                Architecture::Riscv32 | Architecture::Riscv64 => {
                    return Err(String::from("No disassembler for RiscV"))
                }
                Architecture::I386 | Architecture::I586 | Architecture::I686 => Capstone::new()
                    .x86()
                    .mode(arch::x86::ArchMode::Mode32)
                    .build(),
                Architecture::X86_64 => Capstone::new()
                    .x86()
                    .mode(arch::x86::ArchMode::Mode64)
                    .build(),
                Architecture::Arm
                | Architecture::Armv4t
                | Architecture::Armv5te
                | Architecture::Armv7
                | Architecture::Armv7s => Capstone::new().arm().mode(arch::arm::ArchMode::Arm).build(),
                Architecture::Thumbv6m | Architecture::Thumbv7em | Architecture::Thumbv7m => Capstone::new(
                ).arm()
                    .mode(arch::arm::ArchMode::Thumb)
                    .build(),
                Architecture::Aarch64 => Capstone::new()
                    .arm64()
                    .mode(arch::arm64::ArchMode::Arm)
                    .build(),
                _ => return Err(String::from("Unknown ISA")),
            };

            cs.map_err(|err| err.to_string())
        }

        fn print_disassembly(isa: &TargetIsa, mem: &[u8]) -> Result<(), String> {
            let mut cs = get_disassembler(isa)?;

            println!("\nDisassembly of {} bytes:", mem.len());
            let insns = cs.disasm_all(&mem, 0x0).unwrap();
            for i in insns.iter() {
                let mut line = String::new();

                write!(&mut line, "{:4x}:\t", i.address()).unwrap();

                let mut bytes_str = String::new();
                for b in i.bytes() {
                    write!(&mut bytes_str, "{:02x} ", b).unwrap();
                }
                write!(&mut line, "{:21}\t", bytes_str).unwrap();

                if let Some(s) = i.mnemonic() {
                    write!(&mut line, "{}\t", s).unwrap();
                }

                if let Some(s) = i.op_str() {
                    write!(&mut line, "{}", s).unwrap();
                }

                println!("{}", line);
            }
            Ok(())
        }
    } else {
        fn print_disassembly(_: &TargetIsa, _: &[u8]) -> Result<(), String> {
            println!("\nNo disassembly available.");
            Ok(())
        }
    }
}
<commit_msg>Fix copy+pastos in debug messages.<commit_after>//! CLI tool to read Cranelift IR files and compile them into native code.

use cranelift_codegen::isa::TargetIsa;
use cranelift_codegen::print_errors::pretty_error;
use cranelift_codegen::settings::FlagsOrIsa;
use cranelift_codegen::Context;
use cranelift_codegen::{binemit, ir};
use cranelift_reader::parse_test;
use std::path::Path;
use std::path::PathBuf;
use utils::{parse_sets_and_triple, read_to_string};

struct PrintRelocs {
    flag_print: bool,
}

impl binemit::RelocSink for PrintRelocs {
    fn reloc_ebb(
        &mut self,
        where_: binemit::CodeOffset,
        r: binemit::Reloc,
        offset: binemit::CodeOffset,
    ) {
        if self.flag_print {
            println!("reloc_ebb: {} {} at {}", r, offset, where_);
        }
    }

    fn reloc_external(
        &mut self,
        where_: binemit::CodeOffset,
        r: binemit::Reloc,
        name: &ir::ExternalName,
        addend: binemit::Addend,
    ) {
        if self.flag_print {
            println!("reloc_external: {} {} {} at {}", r, name, addend, where_);
        }
    }

    fn reloc_jt(&mut self, where_: binemit::CodeOffset, r: binemit::Reloc, jt: ir::JumpTable) {
        if self.flag_print {
            println!("reloc_jt: {} {} at {}", r, jt, where_);
        }
    }
}

struct PrintTraps {
    flag_print: bool,
}

impl binemit::TrapSink for PrintTraps {
    fn trap(&mut self, offset: binemit::CodeOffset, _srcloc: ir::SourceLoc, code: ir::TrapCode) {
        if self.flag_print {
            println!("trap: {} at {}", code, offset);
        }
    }
}

pub fn run(
    files: Vec<String>,
    flag_print: bool,
    flag_set: &[String],
    flag_isa: &str,
) -> Result<(), String> {
    let parsed = parse_sets_and_triple(flag_set, flag_isa)?;

    for filename in files {
        let path = Path::new(&filename);
        let name = String::from(path.as_os_str().to_string_lossy());
        handle_module(flag_print, &path.to_path_buf(), &name, parsed.as_fisa())?;
    }
    Ok(())
}

fn handle_module(
    flag_print: bool,
    path: &PathBuf,
    name: &str,
    fisa: FlagsOrIsa,
) -> Result<(), String> {
    let buffer = read_to_string(&path).map_err(|e| format!("{}: {}", name, e))?;
    let test_file = parse_test(&buffer, None, None).map_err(|e| format!("{}: {}", name, e))?;

    // If we have an isa from the command-line, use that. Otherwise if the
    // file contains a unique isa, use that.
    let isa = if let Some(isa) = fisa.isa {
        isa
    } else if let Some(isa) = test_file.isa_spec.unique_isa() {
        isa
    } else {
        return Err(String::from("compilation requires a target isa"));
    };

    for (func, _) in test_file.functions {
        let mut context = Context::new();
        context.func = func;

        // Compile and encode the result to machine code.
        let total_size = context
            .compile(isa)
            .map_err(|err| pretty_error(&context.func, Some(isa), err))?;

        let mut mem = Vec::new();
        mem.resize(total_size as usize, 0);

        let mut relocs = PrintRelocs { flag_print };
        let mut traps = PrintTraps { flag_print };
        let mut code_sink: binemit::MemoryCodeSink;
        unsafe {
            code_sink = binemit::MemoryCodeSink::new(mem.as_mut_ptr(), &mut relocs, &mut traps);
        }
        isa.emit_function_to_memory(&context.func, &mut code_sink);

        if flag_print {
            println!("{}", context.func.display(isa));
        }

        if flag_print {
            print!(".byte ");
            let mut first = true;
            for byte in &mem {
                if first {
                    first = false;
                } else {
                    print!(", ");
                }
                print!("{}", byte);
            }

            println!();
            print_disassembly(isa, &mem[0..code_sink.code_size as usize])?;
            print_readonly_data(&mem[code_sink.code_size as usize..total_size as usize]);
        }
    }

    Ok(())
}

fn print_readonly_data(mem: &[u8]) {
    if mem.len() == 0 {
        return;
    }

    println!("\nFollowed by {} bytes of read-only data:", mem.len());

    for (i, byte) in mem.iter().enumerate() {
        if i % 16 == 0 {
            if i != 0 {
                println!();
            }
            print!("{:4}: ", i);
        }
        if i % 4 == 0 {
            print!(" ");
        }
        print!("{:02x} ", byte);
    }
    println!();
}

cfg_if! {
    if #[cfg(feature = "disas")] {
        use capstone::prelude::*;
        use target_lexicon::Architecture;
        use std::fmt::Write;

        fn get_disassembler(isa: &TargetIsa) -> Result<Capstone, String> {
            let cs = match isa.triple().architecture {
                Architecture::Riscv32 | Architecture::Riscv64 => {
                    return Err(String::from("No disassembler for RiscV"))
                }
                Architecture::I386 | Architecture::I586 | Architecture::I686 => Capstone::new()
                    .x86()
                    .mode(arch::x86::ArchMode::Mode32)
                    .build(),
                Architecture::X86_64 => Capstone::new()
                    .x86()
                    .mode(arch::x86::ArchMode::Mode64)
                    .build(),
                Architecture::Arm
                | Architecture::Armv4t
                | Architecture::Armv5te
                | Architecture::Armv7
                | Architecture::Armv7s => Capstone::new().arm().mode(arch::arm::ArchMode::Arm).build(),
                Architecture::Thumbv6m | Architecture::Thumbv7em | Architecture::Thumbv7m => Capstone::new(
                ).arm()
                    .mode(arch::arm::ArchMode::Thumb)
                    .build(),
                Architecture::Aarch64 => Capstone::new()
                    .arm64()
                    .mode(arch::arm64::ArchMode::Arm)
                    .build(),
                _ => return Err(String::from("Unknown ISA")),
            };

            cs.map_err(|err| err.to_string())
        }

        fn print_disassembly(isa: &TargetIsa, mem: &[u8]) -> Result<(), String> {
            let mut cs = get_disassembler(isa)?;

            println!("\nDisassembly of {} bytes:", mem.len());
            let insns = cs.disasm_all(&mem, 0x0).unwrap();
            for i in insns.iter() {
                let mut line = String::new();

                write!(&mut line, "{:4x}:\t", i.address()).unwrap();

                let mut bytes_str = String::new();
                for b in i.bytes() {
                    write!(&mut bytes_str, "{:02x} ", b).unwrap();
                }
                write!(&mut line, "{:21}\t", bytes_str).unwrap();

                if let Some(s) = i.mnemonic() {
                    write!(&mut line, "{}\t", s).unwrap();
                }

                if let Some(s) = i.op_str() {
                    write!(&mut line, "{}", s).unwrap();
                }

                println!("{}", line);
            }
            Ok(())
        }
    } else {
        fn print_disassembly(_: &TargetIsa, _: &[u8]) -> Result<(), String> {
            println!("\nNo disassembly available.");
            Ok(())
        }
    }
}
<|endoftext|>


<commit_before>use super::write_file;
use anyhow::{Context, Result};
use std::path::Path;
use std::{fs, str};

const BINDING_CC_TEMPLATE: &'static str = include_str!("./templates/binding.cc");
const BINDING_GYP_TEMPLATE: &'static str = include_str!("./templates/binding.gyp");
const INDEX_JS_TEMPLATE: &'static str = include_str!("./templates/index.js");
const LIB_RS_TEMPLATE: &'static str = include_str!("./templates/lib.rs");
const BUILD_RS_TEMPLATE: &'static str = include_str!("./templates/build.rs");
const CARGO_TOML_TEMPLATE: &'static str = include_str!("./templates/cargo.toml");
const PACKAGE_JSON_TEMPLATE: &'static str = include_str!("./templates/package.json");
const PARSER_NAME_PLACEHOLDER: &'static str = "PARSER_NAME";
const CLI_VERSION_PLACEHOLDER: &'static str = "CLI_VERSION";
const CLI_VERSION: &'static str = env!("CARGO_PKG_VERSION");

pub fn generate_binding_files(repo_path: &Path, language_name: &str) -> Result<()> {
    let bindings_dir = repo_path.join("bindings");

    let dashed_language_name = language_name.replace("_", "-");
    let dashed_language_name = dashed_language_name.as_str();

    // Generate rust bindings if needed.
    let rust_binding_dir = bindings_dir.join("rust");
    if !rust_binding_dir.exists() {
        create_dir(&rust_binding_dir)?;
        generate_file(
            &rust_binding_dir.join("lib.rs"),
            LIB_RS_TEMPLATE,
            language_name,
        )?;
        generate_file(
            &rust_binding_dir.join("build.rs"),
            BUILD_RS_TEMPLATE,
            language_name,
        )?;
        let cargo_toml_path = repo_path.join("Cargo.toml");
        if !cargo_toml_path.exists() {
            generate_file(&cargo_toml_path, CARGO_TOML_TEMPLATE, dashed_language_name)?;
        }
    }

    // Generate node bindings
    let node_binding_dir = bindings_dir.join("node");
    if !node_binding_dir.exists() {
        create_dir(&node_binding_dir)?;
        generate_file(
            &node_binding_dir.join("index.js"),
            INDEX_JS_TEMPLATE,
            language_name,
        )?;
        generate_file(
            &node_binding_dir.join("binding.cc"),
            BINDING_CC_TEMPLATE,
            language_name,
        )?;

        // Create binding.gyp, or update it with new binding path.
        let binding_gyp_path = repo_path.join("binding.gyp");
        if binding_gyp_path.exists() {
            eprintln!("Updating binding.gyp with new binding path");

            let binding_gyp = fs::read_to_string(&binding_gyp_path)
                .with_context(|| "Failed to read binding.gyp")?;
            let binding_gyp = binding_gyp.replace("src/binding.cc", "bindings/node/binding.cc");
            write_file(&binding_gyp_path, binding_gyp)?;
        } else {
            generate_file(&binding_gyp_path, BINDING_GYP_TEMPLATE, language_name)?;
        }

        // Create package.json, or update it with new binding path.
        let package_json_path = repo_path.join("package.json");
        if package_json_path.exists() {
            let package_json_str = fs::read_to_string(&package_json_path)
                .with_context(|| "Failed to read package.json")?;
            let mut package_json =
                serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(
                    &package_json_str,
                )
                .with_context(|| "Failed to parse package.json")?;
            let package_json_main = package_json.get("main");
            let package_json_needs_update = package_json_main.map_or(true, |v| {
                let main_string = v.as_str();
                main_string == Some("index.js") || main_string == Some("./index.js")
            });
            if package_json_needs_update {
                eprintln!("Updating package.json with new binding path");

                package_json.insert(
                    "main".to_string(),
                    serde_json::Value::String("bindings/node".to_string()),
                );
                let mut package_json_str = serde_json::to_string_pretty(&package_json)?;
                package_json_str.push('\n');
                write_file(&package_json_path, package_json_str)?;
            }
        } else {
            generate_file(
                &package_json_path,
                PACKAGE_JSON_TEMPLATE,
                dashed_language_name,
            )?;
        }

        // Remove files from old node binding paths.
        let old_index_js_path = repo_path.join("index.js");
        let old_binding_cc_path = repo_path.join("src").join("binding.cc");
        if old_index_js_path.exists() {
            fs::remove_file(old_index_js_path).ok();
        }
        if old_binding_cc_path.exists() {
            fs::remove_file(old_binding_cc_path).ok();
        }
    }

    Ok(())
}

fn generate_file(path: &Path, template: &str, language_name: &str) -> Result<()> {
    write_file(
        path,
        template
            .replace(PARSER_NAME_PLACEHOLDER, language_name)
            .replace(CLI_VERSION_PLACEHOLDER, CLI_VERSION),
    )
}

fn create_dir(path: &Path) -> Result<()> {
    fs::create_dir_all(&path)
        .with_context(|| format!("Failed to create {:?}", path.to_string_lossy()))
}
<commit_msg>feat(cli): Independant language binding files generation<commit_after>use super::write_file;
use anyhow::{Context, Result};
use std::path::{Path, PathBuf};
use std::{fs, str};

const BINDING_CC_TEMPLATE: &'static str = include_str!("./templates/binding.cc");
const BINDING_GYP_TEMPLATE: &'static str = include_str!("./templates/binding.gyp");
const INDEX_JS_TEMPLATE: &'static str = include_str!("./templates/index.js");
const LIB_RS_TEMPLATE: &'static str = include_str!("./templates/lib.rs");
const BUILD_RS_TEMPLATE: &'static str = include_str!("./templates/build.rs");
const CARGO_TOML_TEMPLATE: &'static str = include_str!("./templates/cargo.toml");
const PACKAGE_JSON_TEMPLATE: &'static str = include_str!("./templates/package.json");
const PARSER_NAME_PLACEHOLDER: &'static str = "PARSER_NAME";
const CLI_VERSION_PLACEHOLDER: &'static str = "CLI_VERSION";
const CLI_VERSION: &'static str = env!("CARGO_PKG_VERSION");

pub fn generate_binding_files(repo_path: &Path, language_name: &str) -> Result<()> {
    let bindings_dir = repo_path.join("bindings");

    let dashed_language_name = language_name.replace("_", "-");
    let dashed_language_name = dashed_language_name.as_str();

    // Generate rust bindings if needed.
    let rust_binding_dir = bindings_dir.join("rust");
    create_path(&rust_binding_dir, |path| create_dir(path))?;

    create_path(&rust_binding_dir.join("lib.rs").to_owned(), |path| {
        generate_file(path, LIB_RS_TEMPLATE, language_name)
    })?;

    create_path(&rust_binding_dir.join("build.rs").to_owned(), |path| {
        generate_file(path, BUILD_RS_TEMPLATE, language_name)
    })?;

    create_path(&repo_path.join("Cargo.toml").to_owned(), |path| {
        generate_file(path, CARGO_TOML_TEMPLATE, dashed_language_name)
    })?;

    // Generate node bindings
    let node_binding_dir = bindings_dir.join("node");
    let node_binding_dir_created = create_path(&node_binding_dir, |path| create_dir(path))?;

    create_path(&node_binding_dir.join("index.js").to_owned(), |path| {
        generate_file(path, INDEX_JS_TEMPLATE, language_name)
    })?;

    create_path(&node_binding_dir.join("binding.cc").to_owned(), |path| {
        generate_file(path, BINDING_CC_TEMPLATE, language_name)
    })?;

    if node_binding_dir_created {
        // Create binding.gyp, or update it with new binding path.
        let binding_gyp_path = repo_path.join("binding.gyp");
        if binding_gyp_path.exists() {
            eprintln!("Updating binding.gyp with new binding path");

            let binding_gyp = fs::read_to_string(&binding_gyp_path)
                .with_context(|| "Failed to read binding.gyp")?;
            let binding_gyp = binding_gyp.replace("src/binding.cc", "bindings/node/binding.cc");
            write_file(&binding_gyp_path, binding_gyp)?;
        } else {
            generate_file(&binding_gyp_path, BINDING_GYP_TEMPLATE, language_name)?;
        }
    }

    // Create package.json, or update it with new binding path.
    let package_json_path = repo_path.join("package.json");
    if package_json_path.exists() {
        let package_json_str = fs::read_to_string(&package_json_path)
            .with_context(|| "Failed to read package.json")?;
        let mut package_json =
            serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(&package_json_str)
                .with_context(|| "Failed to parse package.json")?;
        let package_json_main = package_json.get("main");
        let package_json_needs_update = package_json_main.map_or(true, |v| {
            let main_string = v.as_str();
            main_string == Some("index.js") || main_string == Some("./index.js")
        });
        if package_json_needs_update {
            eprintln!("Updating package.json with new binding path");

            package_json.insert(
                "main".to_string(),
                serde_json::Value::String("bindings/node".to_string()),
            );
            let mut package_json_str = serde_json::to_string_pretty(&package_json)?;
            package_json_str.push('\n');
            write_file(&package_json_path, package_json_str)?;
        }
    } else {
        generate_file(
            &package_json_path,
            PACKAGE_JSON_TEMPLATE,
            dashed_language_name,
        )?;
    }

    // Remove files from old node binding paths.
    let old_index_js_path = repo_path.join("index.js");
    let old_binding_cc_path = repo_path.join("src").join("binding.cc");
    if old_index_js_path.exists() {
        fs::remove_file(old_index_js_path).ok();
    }
    if old_binding_cc_path.exists() {
        fs::remove_file(old_binding_cc_path).ok();
    }

    Ok(())
}

fn generate_file(path: &Path, template: &str, language_name: &str) -> Result<()> {
    write_file(
        path,
        template
            .replace(PARSER_NAME_PLACEHOLDER, language_name)
            .replace(CLI_VERSION_PLACEHOLDER, CLI_VERSION),
    )
}

fn create_dir(path: &Path) -> Result<()> {
    fs::create_dir_all(&path)
        .with_context(|| format!("Failed to create {:?}", path.to_string_lossy()))
}

fn create_path<F>(path: &PathBuf, action: F) -> Result<bool>
where
    F: Fn(&PathBuf) -> Result<()>,
{
    if !path.exists() {
        action(path)?;
        return Ok(true);
    }
    Ok(false)
}
<|endoftext|>


<commit_before>use std::rc::Rc;

use super::{config, Conditions, Message, TimerEvent};
use state::State;

#[derive(Debug)]
pub enum Context {
    Linear(LinearContext),
}

impl Context {
    pub fn on_timer(&mut self, event: &TimerEvent) {
        match *self {
            Context::Linear(ref mut context) => context.on_timer(event)
        }
    }

    pub fn on_message(&mut self, event: Rc<Message>) {
        match *self {
            Context::Linear(ref mut context) => context.on_message(event)
        }
    }

    pub fn is_open(&self) -> bool {
        match *self {
            Context::Linear(ref context) => context.is_open()
        }
    }

    pub fn new(conditions: Conditions) -> Context {
        Context::Linear(
            LinearContext {
                conditions: conditions,
                state: State::new()
            }
        )
    }
}

#[derive(Debug)]
pub struct LinearContext {
    conditions: Conditions,
    state: State
}

impl LinearContext {
    pub fn on_timer(&mut self, event: &TimerEvent) {
        self.conditions.on_timer(event, &mut self.state);
    }

    pub fn on_message(&mut self, event: Rc<Message>) {
        self.conditions.on_message(event, &mut self.state);
    }

    pub fn is_open(&self) -> bool {
        self.state.is_open()
    }
}

impl From<config::Context> for Context {
    fn from(context: config::Context) -> Context {
        Context::new(context.conditions)
    }
}

use conditions::Builder;

#[test]
fn test_given_close_condition_with_timeout_when_the_timeout_expires_then_the_condition_is_met() {
    let timeout = 100;
    let msg_id = "1".to_string();
    let mut context = Context::new(Builder::new(timeout).patterns(vec![msg_id.clone()]).build());
    let msg1 = btreemap!{
        "uuid".to_string() => msg_id.clone(),
    };
    let event = Rc::new(msg1);
    println!("{:?}", &context);
    assert_false!(context.is_open());
    context.on_message(event);
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(50));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(49));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_false!(context.is_open());
}
#[test]
fn test_given_close_condition_with_max_size_when_the_max_size_reached_then_the_condition_is_met() {
    let timeout = 100;
    let max_size = 3;
    let msg_id = "1".to_string();
    let mut context = Context::new(Builder::new(timeout).max_size(max_size).patterns(vec![msg_id.clone()]).build());
    let msg1 = btreemap!{
        "uuid".to_string() => msg_id.clone(),
    };
    let event = Rc::new(msg1);
    println!("{:?}", &context);
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_message(event.clone());
    println!("{:?}", &context);
    assert_false!(context.is_open());
}

#[test]
fn test_given_close_condition_with_renew_timeout_when_the_timeout_expires_without_renewing_messages_then_the_condition_is_met() {
    let timeout = 100;
    let renew_timeout = 10;
    let msg_id = "1".to_string();
    let mut context = Context::new(Builder::new(timeout).renew_timeout(renew_timeout).patterns(vec![msg_id.clone()]).build());
    let msg1 = btreemap!{
        "uuid".to_string() => msg_id.clone(),
    };
    let event = Rc::new(msg1);
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(8));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_false!(context.is_open());
}

#[test]
fn test_given_close_condition_with_renew_timeout_when_the_timeout_expires_with_renewing_messages_then_the_context_is_not_closed() {
    let timeout = 100;
    let renew_timeout = 10;
    let msg_id = "1".to_string();
    let mut context = Context::new(Builder::new(timeout).renew_timeout(renew_timeout).patterns(vec![msg_id.clone()]).build());
    let msg1 = btreemap!{
        "uuid".to_string() => msg_id.clone(),
    };
    let event = Rc::new(msg1);
    assert_false!(context.is_open());
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(8));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_true!(context.is_open());
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_true!(context.is_open());
}
<commit_msg>create MapContext<commit_after>use std::rc::Rc;

use super::{config, Conditions, Message, TimerEvent};
use state::State;

use self::linear::LinearContext;

#[derive(Debug)]
pub enum Context {
    Linear(LinearContext),
}

impl Context {
    pub fn on_timer(&mut self, event: &TimerEvent) {
        match *self {
            Context::Linear(ref mut context) => context.on_timer(event)
        }
    }

    pub fn on_message(&mut self, event: Rc<Message>) {
        match *self {
            Context::Linear(ref mut context) => context.on_message(event)
        }
    }

    pub fn is_open(&self) -> bool {
        match *self {
            Context::Linear(ref context) => context.is_open()
        }
    }

    pub fn new(conditions: Conditions) -> Context {
        Context::Linear(
            LinearContext::new(conditions)
        )
    }
}

impl From<config::Context> for Context {
    fn from(context: config::Context) -> Context {
        Context::new(context.conditions)
    }
}

mod linear {
    use std::rc::Rc;

    use Conditions;
    use config;
    use Message;
    use state::State;
    use TimerEvent;

    #[derive(Debug)]
    pub struct LinearContext {
        conditions: Conditions,
        state: State
    }

    impl LinearContext {
        pub fn new(conditions: Conditions) -> LinearContext {
            LinearContext {
                conditions: conditions,
                state: State::new()
            }
        }

        pub fn on_timer(&mut self, event: &TimerEvent) {
            self.conditions.on_timer(event, &mut self.state);
        }

        pub fn on_message(&mut self, event: Rc<Message>) {
            self.conditions.on_message(event, &mut self.state);
        }

        pub fn is_open(&self) -> bool {
            self.state.is_open()
        }
    }
}

mod map {
    use std::collections::BTreeMap;
    use std::fmt::Write;
    use std::rc::Rc;

    use Conditions;
    use config;
    use Message;
    use state::State;
    use TimerEvent;

    pub struct MapContext {
        map: BTreeMap<String, State>,
        conditions: Conditions,
        format_buffer: String
    }

    impl MapContext {
        pub fn new(conditions: Conditions) -> MapContext {
            MapContext {
                map: BTreeMap::new(),
                conditions: conditions,
                format_buffer: String::new()
            }
        }

        pub fn on_timer(&mut self, event: &TimerEvent) {
            for (id, mut state) in self.map.iter_mut() {
                self.conditions.on_timer(event, &mut state);
            }
        }

        pub fn on_message(&mut self, event: Rc<Message>) {
            self.format_context_id(&event);
            let state = self.map.entry(self.format_buffer.clone()).or_insert(State::new());
            self.conditions.on_message(event, state);
            self.format_buffer.clear();
        }

        pub fn is_open(&mut self) -> bool {
            !self.map.is_empty()
        }

        fn format_context_id(&mut self, message: &Message) {
            let _ = self.format_buffer.write_str("foo bar baz");
        }
    }
}

use conditions::Builder;

#[test]
fn test_given_close_condition_with_timeout_when_the_timeout_expires_then_the_condition_is_met() {
    let timeout = 100;
    let msg_id = "1".to_string();
    let mut context = Context::new(Builder::new(timeout).patterns(vec![msg_id.clone()]).build());
    let msg1 = btreemap!{
        "uuid".to_string() => msg_id.clone(),
    };
    let event = Rc::new(msg1);
    println!("{:?}", &context);
    assert_false!(context.is_open());
    context.on_message(event);
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(50));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(49));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_false!(context.is_open());
}
#[test]
fn test_given_close_condition_with_max_size_when_the_max_size_reached_then_the_condition_is_met() {
    let timeout = 100;
    let max_size = 3;
    let msg_id = "1".to_string();
    let mut context = Context::new(Builder::new(timeout).max_size(max_size).patterns(vec![msg_id.clone()]).build());
    let msg1 = btreemap!{
        "uuid".to_string() => msg_id.clone(),
    };
    let event = Rc::new(msg1);
    println!("{:?}", &context);
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_message(event.clone());
    println!("{:?}", &context);
    assert_false!(context.is_open());
}

#[test]
fn test_given_close_condition_with_renew_timeout_when_the_timeout_expires_without_renewing_messages_then_the_condition_is_met() {
    let timeout = 100;
    let renew_timeout = 10;
    let msg_id = "1".to_string();
    let mut context = Context::new(Builder::new(timeout).renew_timeout(renew_timeout).patterns(vec![msg_id.clone()]).build());
    let msg1 = btreemap!{
        "uuid".to_string() => msg_id.clone(),
    };
    let event = Rc::new(msg1);
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(8));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_false!(context.is_open());
}

#[test]
fn test_given_close_condition_with_renew_timeout_when_the_timeout_expires_with_renewing_messages_then_the_context_is_not_closed() {
    let timeout = 100;
    let renew_timeout = 10;
    let msg_id = "1".to_string();
    let mut context = Context::new(Builder::new(timeout).renew_timeout(renew_timeout).patterns(vec![msg_id.clone()]).build());
    let msg1 = btreemap!{
        "uuid".to_string() => msg_id.clone(),
    };
    let event = Rc::new(msg1);
    assert_false!(context.is_open());
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(8));
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_true!(context.is_open());
    context.on_message(event.clone());
    assert_true!(context.is_open());
    context.on_timer(&mut TimerEvent(1));
    assert_true!(context.is_open());
}
<|endoftext|>