<commit_before>extern crate c_ares_sys;
extern crate libc;

use std::ffi::CString;
use std::mem;
use std::os::unix::io;
use std::ptr;

use callbacks::{
    socket_callback,
    query_a_callback,
    query_aaaa_callback,
};
use types::{
    Flag,
    AresError,
    AResult,
    AAAAResult,
    DnsClass,
    QueryType,
};
use utils::ares_error;

/// Used to configure the behaviour of the name resolver.
pub struct Options {
    ares_options: c_ares_sys::Struct_ares_options,
    optmask: libc::c_int,
    domains: Vec<CString>,
    lookups: Option<CString>,
}

impl Options {
    /// Returns a fresh `Options`, on which no values are set.
    pub fn new() -> Options {
        Options {
            ares_options: c_ares_sys::Struct_ares_options::default(),
            optmask: 0,
            domains: Vec::new(),
            lookups: None,
        }
    }

    /// Set flags controlling the behaviour of the resolver.
    pub fn set_flags(&mut self, flags: &[Flag]) -> &mut Self {
        let c_flags: libc::c_int = flags
            .iter()
            .fold(0, |acc, &flag| acc | (flag as libc::c_int));
        self.ares_options.flags = c_flags;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_FLAGS;
        self
    }

    /// Set the number of milliseconds each name server is given to respond to
    /// a query on the first try.  (After the first try, the timeout algorithm
    /// becomes more complicated, but scales linearly with the value of
    /// timeout.) The default is five seconds.
    pub fn set_timeout(&mut self, ms: u32) -> &mut Self {
        self.ares_options.timeout = ms as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_TIMEOUTMS;
        self
    }

    /// Set the number of tries the resolver will try contacting each name
    /// server before giving up. The default is four tries.
    pub fn set_tries(&mut self, tries: u32) -> &mut Self {
        self.ares_options.tries = tries as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_TRIES;
        self
    }

    /// Set the number of dots which must be present in a domain name for it to
    /// be queried for "as is" prior to querying for it with the default domain
    /// extensions appended. The default value is 1 unless set otherwise by
    /// resolv.conf or the RES_OPTIONS environment variable.
    pub fn set_ndots(&mut self, ndots: u32) -> &mut Self {
        self.ares_options.ndots = ndots as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_NDOTS;
        self
    }

    /// Set the UDP port to use for queries. The default value is 53, the
    /// standard name service port.
    pub fn set_udp_port(&mut self, udp_port: u16) -> &mut Self {
        self.ares_options.udp_port = udp_port as libc::c_ushort;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_UDP_PORT;
        self
    }

    /// Set the TCP port to use for queries. The default value is 53, the
    /// standard name service port.
    pub fn set_tcp_port(&mut self, tcp_port: u16) -> &mut Self {
        self.ares_options.tcp_port = tcp_port as libc::c_ushort;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_TCP_PORT;
        self
    }

    /// Set the domains to search, instead of the domains specified in
    /// resolv.conf or the domain derived from the kernel hostname variable.
    pub fn set_domains(&mut self, domains: &[&str]) -> &mut Self {
        self.domains = domains
            .iter()
            .map(|&s| CString::new(s).unwrap())
            .collect();
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_DOMAINS;
        self
    }

    /// Set the lookups to perform for host queries. `lookups` should be set to
    /// a string of the characters "b" or "f", where "b" indicates a DNS lookup
    /// and "f" indicates a lookup in the hosts file.
    pub fn set_lookups(&mut self, lookups: &str) -> &mut Self {
        let c_lookups = CString::new(lookups).unwrap();
        self.lookups = Some(c_lookups);
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_LOOKUPS;
        self
    }

    /// Set the socket send buffer size.
    pub fn set_sock_send_buffer_size(&mut self, size: u32) -> &mut Self {
        self.ares_options.socket_send_buffer_size = size as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_SOCK_SNDBUF;
        self
    }

    /// Set the socket receive buffer size.
    pub fn set_sock_receive_buffer_size(&mut self, size: u32) -> &mut Self {
        self.ares_options.socket_receive_buffer_size = size as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_SOCK_RCVBUF;
        self
    }

    /// Set the EDNS packet size.
    pub fn set_ednspsz(&mut self, size: u32) -> &mut Self {
        self.ares_options.ednspsz = size as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_EDNSPSZ;
        self
    }
}

/// A channel for name service lookups.
pub struct Channel {
    ares_channel: c_ares_sys::ares_channel,
}

impl Channel {
    /// Create a new channel for name service lookups, providing a callback
    /// for socket state changes.
    ///
    /// `callback(socket, read, write)` will be called when a socket changes
    /// state:
    ///
    /// -  `read` is set to true if the socket should listen for read events
    /// -  `write` is set to true if the socket should listen to write events.
    pub fn new<F>(
        callback: F,
        mut options: Options) -> Result<Channel, AresError>
        where F: FnMut(io::RawFd, bool, bool) + 'static {
        options.optmask = options.optmask | c_ares_sys::ARES_OPT_SOCK_STATE_CB;
        options.ares_options.sock_state_cb = Some(socket_callback::<F>);
        options.ares_options.sock_state_cb_data = unsafe {
            mem::transmute(Box::new(callback))
        };
        Self::create_channel(options)
    }

    fn create_channel(mut options: Options) -> Result<Channel, AresError> {
        // Initialize the library.
        let lib_rc = unsafe {
            c_ares_sys::ares_library_init(c_ares_sys::ARES_LIB_INIT_ALL)
        };
        if lib_rc != c_ares_sys::ARES_SUCCESS {
            return Err(ares_error(lib_rc))
        }

        // We deferred setting up domains in the options - do it now.
        let domains: Vec<_> = options.domains
            .iter()
            .map(|s| s.as_ptr())
            .collect();
        options.ares_options.domains =
            domains.as_ptr() as *mut *mut libc::c_char;
        options.ares_options.ndomains = domains.len() as libc::c_int;

        // Likewise for lookups.
        for c_lookup in options.lookups.iter() {
            options.ares_options.lookups =
                c_lookup.as_ptr() as *mut libc::c_char;
        }

        // Initialize the channel.
        let mut ares_channel = ptr::null_mut();
        let channel_rc = unsafe {
            c_ares_sys::ares_init_options(
                &mut ares_channel,
                &mut options.ares_options,
                options.optmask)
        };
        if channel_rc != c_ares_sys::ARES_SUCCESS {
            unsafe { c_ares_sys::ares_library_cleanup(); }
            return Err(ares_error(channel_rc))
        }

        let channel = Channel {
            ares_channel: ares_channel,
        };
        Ok(channel)
    }

    /// Handle input, output, and timeout events associated with the specified
    /// file descriptors (sockets).
    ///
    /// Providing a value for `read_fd` indicates that the identified socket
    /// is readable; likewise providing a value for `write_fd` indicates that
    /// the identified socket is writable.  Use `INVALID_FD` for "no action".
    pub fn process_fd(&mut self, read_fd: io::RawFd, write_fd: io::RawFd) {
        unsafe {
            c_ares_sys::ares_process_fd(
                self.ares_channel,
                read_fd as c_ares_sys::ares_socket_t,
                write_fd as c_ares_sys::ares_socket_t);
        }
    }

    /// Set the list of servers to contact, instead of the servers specified
    /// in resolv.conf or the local named.
    ///
    /// String format is `host[:port]`.  IPv6 addresses with ports require
    /// square brackets eg `[2001:4860:4860::8888]:53`.
    pub fn set_servers(&mut self, servers: &[&str]) -> Result<(), AresError> {
        let servers_csv = servers.connect(",");
        let c_servers = CString::new(servers_csv).unwrap();
        let ares_rc = unsafe {
            c_ares_sys::ares_set_servers_csv(
                self.ares_channel,
                c_servers.as_ptr())
        };
        if ares_rc != c_ares_sys::ARES_SUCCESS {
            Err(ares_error(ares_rc))
        } else {
            Ok(())
        }
    }

    /// Look up the A record associated with `name`.
    ///
    /// On completion, `handler` is called with the result.
    pub fn query_a<F>(&mut self, name: &str, handler: F)
        where F: FnOnce(Result<AResult, AresError>) + 'static {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let c_arg: *mut libc::c_void = mem::transmute(Box::new(handler));
            c_ares_sys::ares_query(
                self.ares_channel,
                c_name.as_ptr(),
                DnsClass::IN as libc::c_int,
                QueryType::A as libc::c_int,
                Some(query_a_callback::<F>),
                c_arg);
        }
    }

    /// Look up the AAAA record associated with `name`.
    ///
    /// On completion, `handler` is called with the result.
    pub fn query_aaaa<F>(&mut self, name: &str, handler: F)
        where F: FnOnce(Result<AAAAResult, AresError>) + 'static {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let c_arg: *mut libc::c_void = mem::transmute(Box::new(handler));
            c_ares_sys::ares_query(
                self.ares_channel,
                c_name.as_ptr(),
                DnsClass::IN as libc::c_int,
                QueryType::AAAA as libc::c_int,
                Some(query_aaaa_callback::<F>),
                c_arg);
        }
    }
}

impl Drop for Channel {
    fn drop(&mut self) {
        unsafe {
            c_ares_sys::ares_destroy(self.ares_channel);
            c_ares_sys::ares_library_cleanup();
        }
    }
}

unsafe impl Send for Channel { }
<commit_msg>Add support for 'rotate' option<commit_after>extern crate c_ares_sys;
extern crate libc;

use std::ffi::CString;
use std::mem;
use std::os::unix::io;
use std::ptr;

use callbacks::{
    socket_callback,
    query_a_callback,
    query_aaaa_callback,
};
use types::{
    Flag,
    AresError,
    AResult,
    AAAAResult,
    DnsClass,
    QueryType,
};
use utils::ares_error;

/// Used to configure the behaviour of the name resolver.
pub struct Options {
    ares_options: c_ares_sys::Struct_ares_options,
    optmask: libc::c_int,
    domains: Vec<CString>,
    lookups: Option<CString>,
}

impl Options {
    /// Returns a fresh `Options`, on which no values are set.
    pub fn new() -> Options {
        Options {
            ares_options: c_ares_sys::Struct_ares_options::default(),
            optmask: 0,
            domains: Vec::new(),
            lookups: None,
        }
    }

    /// Set flags controlling the behaviour of the resolver.
    pub fn set_flags(&mut self, flags: &[Flag]) -> &mut Self {
        let c_flags: libc::c_int = flags
            .iter()
            .fold(0, |acc, &flag| acc | (flag as libc::c_int));
        self.ares_options.flags = c_flags;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_FLAGS;
        self
    }

    /// Set the number of milliseconds each name server is given to respond to
    /// a query on the first try.  (After the first try, the timeout algorithm
    /// becomes more complicated, but scales linearly with the value of
    /// timeout.) The default is five seconds.
    pub fn set_timeout(&mut self, ms: u32) -> &mut Self {
        self.ares_options.timeout = ms as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_TIMEOUTMS;
        self
    }

    /// Set the number of tries the resolver will try contacting each name
    /// server before giving up. The default is four tries.
    pub fn set_tries(&mut self, tries: u32) -> &mut Self {
        self.ares_options.tries = tries as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_TRIES;
        self
    }

    /// Set the number of dots which must be present in a domain name for it to
    /// be queried for "as is" prior to querying for it with the default domain
    /// extensions appended. The default value is 1 unless set otherwise by
    /// resolv.conf or the RES_OPTIONS environment variable.
    pub fn set_ndots(&mut self, ndots: u32) -> &mut Self {
        self.ares_options.ndots = ndots as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_NDOTS;
        self
    }

    /// Set the UDP port to use for queries. The default value is 53, the
    /// standard name service port.
    pub fn set_udp_port(&mut self, udp_port: u16) -> &mut Self {
        self.ares_options.udp_port = udp_port as libc::c_ushort;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_UDP_PORT;
        self
    }

    /// Set the TCP port to use for queries. The default value is 53, the
    /// standard name service port.
    pub fn set_tcp_port(&mut self, tcp_port: u16) -> &mut Self {
        self.ares_options.tcp_port = tcp_port as libc::c_ushort;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_TCP_PORT;
        self
    }

    /// Set the domains to search, instead of the domains specified in
    /// resolv.conf or the domain derived from the kernel hostname variable.
    pub fn set_domains(&mut self, domains: &[&str]) -> &mut Self {
        self.domains = domains
            .iter()
            .map(|&s| CString::new(s).unwrap())
            .collect();
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_DOMAINS;
        self
    }

    /// Set the lookups to perform for host queries. `lookups` should be set to
    /// a string of the characters "b" or "f", where "b" indicates a DNS lookup
    /// and "f" indicates a lookup in the hosts file.
    pub fn set_lookups(&mut self, lookups: &str) -> &mut Self {
        let c_lookups = CString::new(lookups).unwrap();
        self.lookups = Some(c_lookups);
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_LOOKUPS;
        self
    }

    /// Set the socket send buffer size.
    pub fn set_sock_send_buffer_size(&mut self, size: u32) -> &mut Self {
        self.ares_options.socket_send_buffer_size = size as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_SOCK_SNDBUF;
        self
    }

    /// Set the socket receive buffer size.
    pub fn set_sock_receive_buffer_size(&mut self, size: u32) -> &mut Self {
        self.ares_options.socket_receive_buffer_size = size as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_SOCK_RCVBUF;
        self
    }

    /// Configure round robin selection of nameservers.
    pub fn set_rotate(&mut self) -> &mut Self {
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_ROTATE;
        self
    }

    /// Set the EDNS packet size.
    pub fn set_ednspsz(&mut self, size: u32) -> &mut Self {
        self.ares_options.ednspsz = size as libc::c_int;
        self.optmask = self.optmask | c_ares_sys::ARES_OPT_EDNSPSZ;
        self
    }
}

/// A channel for name service lookups.
pub struct Channel {
    ares_channel: c_ares_sys::ares_channel,
}

impl Channel {
    /// Create a new channel for name service lookups, providing a callback
    /// for socket state changes.
    ///
    /// `callback(socket, read, write)` will be called when a socket changes
    /// state:
    ///
    /// -  `read` is set to true if the socket should listen for read events
    /// -  `write` is set to true if the socket should listen to write events.
    pub fn new<F>(
        callback: F,
        mut options: Options) -> Result<Channel, AresError>
        where F: FnMut(io::RawFd, bool, bool) + 'static {
        options.optmask = options.optmask | c_ares_sys::ARES_OPT_SOCK_STATE_CB;
        options.ares_options.sock_state_cb = Some(socket_callback::<F>);
        options.ares_options.sock_state_cb_data = unsafe {
            mem::transmute(Box::new(callback))
        };
        Self::create_channel(options)
    }

    fn create_channel(mut options: Options) -> Result<Channel, AresError> {
        // Initialize the library.
        let lib_rc = unsafe {
            c_ares_sys::ares_library_init(c_ares_sys::ARES_LIB_INIT_ALL)
        };
        if lib_rc != c_ares_sys::ARES_SUCCESS {
            return Err(ares_error(lib_rc))
        }

        // We deferred setting up domains in the options - do it now.
        let domains: Vec<_> = options.domains
            .iter()
            .map(|s| s.as_ptr())
            .collect();
        options.ares_options.domains =
            domains.as_ptr() as *mut *mut libc::c_char;
        options.ares_options.ndomains = domains.len() as libc::c_int;

        // Likewise for lookups.
        for c_lookup in options.lookups.iter() {
            options.ares_options.lookups =
                c_lookup.as_ptr() as *mut libc::c_char;
        }

        // Initialize the channel.
        let mut ares_channel = ptr::null_mut();
        let channel_rc = unsafe {
            c_ares_sys::ares_init_options(
                &mut ares_channel,
                &mut options.ares_options,
                options.optmask)
        };
        if channel_rc != c_ares_sys::ARES_SUCCESS {
            unsafe { c_ares_sys::ares_library_cleanup(); }
            return Err(ares_error(channel_rc))
        }

        let channel = Channel {
            ares_channel: ares_channel,
        };
        Ok(channel)
    }

    /// Handle input, output, and timeout events associated with the specified
    /// file descriptors (sockets).
    ///
    /// Providing a value for `read_fd` indicates that the identified socket
    /// is readable; likewise providing a value for `write_fd` indicates that
    /// the identified socket is writable.  Use `INVALID_FD` for "no action".
    pub fn process_fd(&mut self, read_fd: io::RawFd, write_fd: io::RawFd) {
        unsafe {
            c_ares_sys::ares_process_fd(
                self.ares_channel,
                read_fd as c_ares_sys::ares_socket_t,
                write_fd as c_ares_sys::ares_socket_t);
        }
    }

    /// Set the list of servers to contact, instead of the servers specified
    /// in resolv.conf or the local named.
    ///
    /// String format is `host[:port]`.  IPv6 addresses with ports require
    /// square brackets eg `[2001:4860:4860::8888]:53`.
    pub fn set_servers(&mut self, servers: &[&str]) -> Result<(), AresError> {
        let servers_csv = servers.connect(",");
        let c_servers = CString::new(servers_csv).unwrap();
        let ares_rc = unsafe {
            c_ares_sys::ares_set_servers_csv(
                self.ares_channel,
                c_servers.as_ptr())
        };
        if ares_rc != c_ares_sys::ARES_SUCCESS {
            Err(ares_error(ares_rc))
        } else {
            Ok(())
        }
    }

    /// Look up the A record associated with `name`.
    ///
    /// On completion, `handler` is called with the result.
    pub fn query_a<F>(&mut self, name: &str, handler: F)
        where F: FnOnce(Result<AResult, AresError>) + 'static {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let c_arg: *mut libc::c_void = mem::transmute(Box::new(handler));
            c_ares_sys::ares_query(
                self.ares_channel,
                c_name.as_ptr(),
                DnsClass::IN as libc::c_int,
                QueryType::A as libc::c_int,
                Some(query_a_callback::<F>),
                c_arg);
        }
    }

    /// Look up the AAAA record associated with `name`.
    ///
    /// On completion, `handler` is called with the result.
    pub fn query_aaaa<F>(&mut self, name: &str, handler: F)
        where F: FnOnce(Result<AAAAResult, AresError>) + 'static {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let c_arg: *mut libc::c_void = mem::transmute(Box::new(handler));
            c_ares_sys::ares_query(
                self.ares_channel,
                c_name.as_ptr(),
                DnsClass::IN as libc::c_int,
                QueryType::AAAA as libc::c_int,
                Some(query_aaaa_callback::<F>),
                c_arg);
        }
    }
}

impl Drop for Channel {
    fn drop(&mut self) {
        unsafe {
            c_ares_sys::ares_destroy(self.ares_channel);
            c_ares_sys::ares_library_cleanup();
        }
    }
}

unsafe impl Send for Channel { }
<|endoftext|>


<commit_before>#![feature(core_intrinsics)]
#![feature(stmt_expr_attributes)]

extern crate rand;
extern crate errno;
#[cfg(unix)] extern crate libc;
#[cfg(windows)] extern crate kernel32;
#[cfg(all(unix, test))] extern crate nix;

mod alloc;

pub use alloc::{ unprotected_mprotect, malloc, allocarray, free };

// -- memcmp --

/// Constant time memcmp.
pub unsafe fn memcmp<T>(b1: *const T, b2: *const T, len: usize) -> bool {
    let b1 = b1 as *const u8;
    let b2 = b2 as *const u8;
    let mut d = 0;
    for i in 0..len as isize {
        d |= *b1.offset(i) ^ *b2.offset(i);
    }
    d == 0
}

// -- memset / memzero --

/// General memset.
#[cfg(not(HAVE_MEMSET_S))]
pub unsafe fn memset<T>(s: *mut T, c: i32, n: usize) {
    let s = s as *mut u8;
    let mut i = 0;
    while i < n {
        *s.offset(i as isize) = c as u8;
        i += 1;
    }
}

/// Call memset_s.
#[cfg(HAVE_MEMSET_S)]
pub unsafe fn memset<T>(s: *mut T, c: i32, n: usize) {
    extern {
        fn memset_s(s: *mut libc::c_void, smax: libc::size_t, c: libc::c_int, n: libc::size_t) -> libc::c_int;
    }
    memset_s(s as libc::c_void, n, c, n);
}

/// General memzero.
#[cfg(all(unix, not(HAVE_EXPLICIT_BZERO)))]
pub unsafe fn memzero<T>(dest: *mut T, n: usize) {
    memset(dest, 0, n);
}

/// Call explicit_bzero.
#[cfg(all(unix, HAVE_EXPLICIT_BZERO))]
pub unsafe fn memzero<T>(dest: *mut T, n: usize) {
    extern {
        fn explicit_bzero(s: *mut libc::c_void, n: libc::size_t);
    }
    explicit_bzero(dest as *mut libc::c_void, n);
}

/// call SecureZeroMemory.
#[cfg(windows)]
pub unsafe fn memzero<T>(s: *mut T, n: usize) {
    extern "system" {
        fn SecureZeroMemory(ptr: winapi::PVOID, cnt: winapi::SIZE_T);
    }
    SecureZeroMemory(s as winapi::PVOID, n as winapi::SIZE_T);
}

// -- mlock / munlock --

/// Unix mlock.
#[cfg(unix)]
pub unsafe fn mlock<T>(addr: *mut T, len: usize) -> bool {
    libc::madvise(addr as *mut libc::c_void, len, libc::MADV_DONTDUMP);
    libc::mlock(addr as *mut libc::c_void, len) == 0
}

/// Windows VirtualLock.
#[cfg(windows)]
pub unsafe fn mlock<T>(addr: *mut T, len: usize) -> bool {
    kernel32::VirtualLock(addr as winapi::LPVOID, len as winapi::SIZE_T) != 0
}

/// Unix munlock.
#[cfg(unix)]
pub unsafe fn munlock<T>(addr: *mut T, len: usize) -> bool {
    memzero(addr, len);
    libc::madvise(addr as *mut libc::c_void, len, libc::MADV_DODUMP);
    libc::munlock(addr as *mut libc::c_void, len) == 0
}

/// Windows VirtualUnlock.
#[cfg(windows)]
pub unsafe fn munlock<T>(addr: *mut T, len: usize) -> bool {
    memzero(addr, len);
    kernel32::VirtualUnlock(addr as winapi::LPVOID, len as winapi::SIZE_T) != 0
}

// -- mprotect --

/// Prot enum.
#[derive(Debug, Clone, PartialEq)]
pub enum Prot {
    #[cfg(unix)] NoAccess = libc::PROT_NONE as isize,
    #[cfg(unix)] ReadOnly = libc::PROT_READ as isize,
    #[cfg(unix)] WriteOnly = libc::PROT_WRITE as isize,
    #[cfg(unix)] ReadWrite = (libc::PROT_READ | libc::PROT_WRITE) as isize,
    #[cfg(unix)] Execute = libc::PROT_EXEC as isize,
    #[cfg(unix)] ReadExec = (libc::PROT_READ | libc::PROT_EXEC) as isize,
    #[cfg(unix)] WriteExec = (libc::PROT_WRITE | libc::PROT_EXEC) as isize,
    #[cfg(unix)] ReadWriteExec = (libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC) as isize,

    #[cfg(windows)] NoAccess = winapi::PAGE_NOACCESS as isize,
    #[cfg(windows)] ReadOnly = winapi::PAGE_READONLY as isize,
    #[cfg(windows)] ReadWrite = winapi::PAGE_READWRITE as isize,
    #[cfg(windows)] WriteCopy = winapi::PAGE_WRITECOPY as isize,
    #[cfg(windows)] Execute = winapi::PAGE_EXECUTE as isize,
    #[cfg(windows)] ReadExec = winapi::PAGE_EXECUTE_READ as isize,
    #[cfg(windows)] ReadWriteExec = winapi::PAGE_EXECUTE_READWRITE as isize,
    #[cfg(windows)] WriteCopyExec = winapi::PAGE_EXECUTE_WRITECOPY as isize,
    #[cfg(windows)] Guard = winapi::PAGE_GUARD as isize,
    #[cfg(windows)] NoCache = winapi::PAGE_NOCACHE as isize,
    #[cfg(windows)] WriteCombine = winapi::PAGE_WRITECOMBINE as isize,
    #[cfg(windows)] RevertToFileMap = winapi::PAGE_REVERT_TO_FILE_MAP as isize,
    #[cfg(windows)] TargetsNoUpdate = winapi::PAGE_TARGETS_NO_UPDATE as isize,
    #[cfg(windows)] TargetsInvalid = winapi::PAGE_TARGETS_INVALID as isize,
}

/// Unix mprotect.
#[cfg(unix)]
pub unsafe fn mprotect<T>(ptr: *mut T, len: usize, prot: Prot) -> bool {
    libc::mprotect(ptr as *mut libc::c_void, len, prot as libc::c_int) == 0
}

/// Windows VirtualProtect.
#[cfg(windows)]
pub unsafe fn mprotect<T>(ptr: *mut T, len: usize, prot: Prot) -> bool {
    kernel32::VirtualProtect(
        ptr as winapi::LPVOID,
        len as winapi::SIZE_T,
        prot as winapi::DWORD,
        core::mem::uninitialized()
    ) != 0
}
<commit_msg>MADV_DONTDUMP is not universally available<commit_after>#![feature(core_intrinsics)]
#![feature(stmt_expr_attributes)]

extern crate rand;
extern crate errno;
#[cfg(unix)] extern crate libc;
#[cfg(windows)] extern crate kernel32;
#[cfg(all(unix, test))] extern crate nix;

mod alloc;

pub use alloc::{ unprotected_mprotect, malloc, allocarray, free };

// -- memcmp --

/// Constant time memcmp.
pub unsafe fn memcmp<T>(b1: *const T, b2: *const T, len: usize) -> bool {
    let b1 = b1 as *const u8;
    let b2 = b2 as *const u8;
    let mut d = 0;
    for i in 0..len as isize {
        d |= *b1.offset(i) ^ *b2.offset(i);
    }
    d == 0
}

// -- memset / memzero --

/// General memset.
#[cfg(not(HAVE_MEMSET_S))]
pub unsafe fn memset<T>(s: *mut T, c: i32, n: usize) {
    let s = s as *mut u8;
    let mut i = 0;
    while i < n {
        *s.offset(i as isize) = c as u8;
        i += 1;
    }
}

/// Call memset_s.
#[cfg(HAVE_MEMSET_S)]
pub unsafe fn memset<T>(s: *mut T, c: i32, n: usize) {
    extern {
        fn memset_s(s: *mut libc::c_void, smax: libc::size_t, c: libc::c_int, n: libc::size_t) -> libc::c_int;
    }
    memset_s(s as libc::c_void, n, c, n);
}

/// General memzero.
#[cfg(all(unix, not(HAVE_EXPLICIT_BZERO)))]
pub unsafe fn memzero<T>(dest: *mut T, n: usize) {
    memset(dest, 0, n);
}

/// Call explicit_bzero.
#[cfg(all(unix, HAVE_EXPLICIT_BZERO))]
pub unsafe fn memzero<T>(dest: *mut T, n: usize) {
    extern {
        fn explicit_bzero(s: *mut libc::c_void, n: libc::size_t);
    }
    explicit_bzero(dest as *mut libc::c_void, n);
}

/// call SecureZeroMemory.
#[cfg(windows)]
pub unsafe fn memzero<T>(s: *mut T, n: usize) {
    extern "system" {
        fn SecureZeroMemory(ptr: winapi::PVOID, cnt: winapi::SIZE_T);
    }
    SecureZeroMemory(s as winapi::PVOID, n as winapi::SIZE_T);
}

// -- mlock / munlock --

#[cfg(target_os = "linux")]
unsafe fn dontdump<T>(addr: *mut T, len: usize) {
    libc::madvise(addr as *mut libc::c_void, len, libc::MADV_DONTDUMP);
}

#[cfg(target_os = "linux")]
unsafe fn dodump<T>(addr: *mut T, len: usize) {
    libc::madvise(addr as *mut libc::c_void, len, libc::MADV_DODUMP);
}

#[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
unsafe fn dontdump<T>(addr: *mut T, len: usize) {
    libc::madvise(addr as *mut libc::c_void, len, libc::MADV_NOCORE);
}

#[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
unsafe fn dodump<T>(addr: *mut T, len: usize) {
    libc::madvise(addr as *mut libc::c_void, len, libc::MADV_CORE);
}

#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "dragonfly")))]
fn dontdump<T>(_addr: *mut T, _len: usize) { }

#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "dragonfly")))]
fn dodump<T>(_addr: *mut T, _len: usize) { }

/// Unix mlock.
#[cfg(unix)]
pub unsafe fn mlock<T>(addr: *mut T, len: usize) -> bool {
    dontdump(addr, len);
    libc::mlock(addr as *mut libc::c_void, len) == 0
}

/// Windows VirtualLock.
#[cfg(windows)]
pub unsafe fn mlock<T>(addr: *mut T, len: usize) -> bool {
    kernel32::VirtualLock(addr as winapi::LPVOID, len as winapi::SIZE_T) != 0
}

/// Unix munlock.
#[cfg(unix)]
pub unsafe fn munlock<T>(addr: *mut T, len: usize) -> bool {
    memzero(addr, len);
    dodump(addr, len);
    libc::munlock(addr as *mut libc::c_void, len) == 0
}

/// Windows VirtualUnlock.
#[cfg(windows)]
pub unsafe fn munlock<T>(addr: *mut T, len: usize) -> bool {
    memzero(addr, len);
    kernel32::VirtualUnlock(addr as winapi::LPVOID, len as winapi::SIZE_T) != 0
}

// -- mprotect --

/// Prot enum.
#[derive(Debug, Clone, PartialEq)]
pub enum Prot {
    #[cfg(unix)] NoAccess = libc::PROT_NONE as isize,
    #[cfg(unix)] ReadOnly = libc::PROT_READ as isize,
    #[cfg(unix)] WriteOnly = libc::PROT_WRITE as isize,
    #[cfg(unix)] ReadWrite = (libc::PROT_READ | libc::PROT_WRITE) as isize,
    #[cfg(unix)] Execute = libc::PROT_EXEC as isize,
    #[cfg(unix)] ReadExec = (libc::PROT_READ | libc::PROT_EXEC) as isize,
    #[cfg(unix)] WriteExec = (libc::PROT_WRITE | libc::PROT_EXEC) as isize,
    #[cfg(unix)] ReadWriteExec = (libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC) as isize,

    #[cfg(windows)] NoAccess = winapi::PAGE_NOACCESS as isize,
    #[cfg(windows)] ReadOnly = winapi::PAGE_READONLY as isize,
    #[cfg(windows)] ReadWrite = winapi::PAGE_READWRITE as isize,
    #[cfg(windows)] WriteCopy = winapi::PAGE_WRITECOPY as isize,
    #[cfg(windows)] Execute = winapi::PAGE_EXECUTE as isize,
    #[cfg(windows)] ReadExec = winapi::PAGE_EXECUTE_READ as isize,
    #[cfg(windows)] ReadWriteExec = winapi::PAGE_EXECUTE_READWRITE as isize,
    #[cfg(windows)] WriteCopyExec = winapi::PAGE_EXECUTE_WRITECOPY as isize,
    #[cfg(windows)] Guard = winapi::PAGE_GUARD as isize,
    #[cfg(windows)] NoCache = winapi::PAGE_NOCACHE as isize,
    #[cfg(windows)] WriteCombine = winapi::PAGE_WRITECOMBINE as isize,
    #[cfg(windows)] RevertToFileMap = winapi::PAGE_REVERT_TO_FILE_MAP as isize,
    #[cfg(windows)] TargetsNoUpdate = winapi::PAGE_TARGETS_NO_UPDATE as isize,
    #[cfg(windows)] TargetsInvalid = winapi::PAGE_TARGETS_INVALID as isize,
}

/// Unix mprotect.
#[cfg(unix)]
pub unsafe fn mprotect<T>(ptr: *mut T, len: usize, prot: Prot) -> bool {
    libc::mprotect(ptr as *mut libc::c_void, len, prot as libc::c_int) == 0
}

/// Windows VirtualProtect.
#[cfg(windows)]
pub unsafe fn mprotect<T>(ptr: *mut T, len: usize, prot: Prot) -> bool {
    kernel32::VirtualProtect(
        ptr as winapi::LPVOID,
        len as winapi::SIZE_T,
        prot as winapi::DWORD,
        core::mem::uninitialized()
    ) != 0
}
<|endoftext|>


<commit_before>use regex::Regex;
use std::sync::mpsc::channel;

use adapter::ChatAdapter;
use handler::MessageHandler;

/// The Chatbot is the central data structure of the chatbot platform. It contains a `run` method
/// which listens for messages from adapters and routes them to handlers. Any program which uses
/// chatbot will need to minimally create a Chatbot, add an adapter, add a handler, and call Chatbot
/// [`run`](chatbot/struct.Chatbot.html#method.run).
pub struct Chatbot {
    name: String,
    adapters: Vec<Box<ChatAdapter>>,
    handlers: Vec<Box<MessageHandler>>,
    addresser: Option<Regex>
}

impl Chatbot {
    /// Create a new chatbot instance.
    ///
    /// The name provided here will be used for message filtering and in handlers for whatever they
    /// want. You'll want to make your binding mutable so you can call `add_adapter` and
    /// `add_handler`.
    pub fn new(name: &str) -> Chatbot {
        Chatbot {
            name: name.to_owned(),
            adapters: Vec::new(),
            handlers: Vec::new(),
            addresser: None
        }
    }

    /// Return the name provided in initialization
    pub fn get_name(&self) -> &str {
        self.name.as_ref()
    }

    pub fn get_addresser(&self) -> Option<&Regex> {
        self.addresser.as_ref()
    }

    pub fn address_by_name(&mut self) {
        let addresser_str = format!(r"^\s*@?{}[:,\s]\s*", self.get_name());
        let addresser = Regex::new(addresser_str.as_ref());

        self.addresser = addresser.ok();
    }

    /// Add a ChatAdapter to the bot
    ///
    /// Add as many adapters as you like. The IncomingMessages sent by adapters are made available
    /// to all handlers regardless of how many adapters exist. The IncomingMessage.reply method
    /// makes sure the response is sent back to the adapter from whence the message came.
    pub fn add_adapter<T>(&mut self, adapter: T)
        where T: ChatAdapter + 'static
    {
        println!("Adding adapter {}", adapter.get_name());
        self.adapters.push(Box::new(adapter))
    }

    /// Add a MessageHandler to the bot
    ///
    /// The more handlers you have the more useful your bot becomes (for potentially loose
    /// definitions of useful :P). Check out the handler! macro for making simple handlers and also
    /// see the `MessageHandler` implementors section for a list of built in handlers.
    pub fn add_handler<T>(&mut self, handler: T)
        where T: MessageHandler + 'static
    {
        println!("Adding handler {}", handler.name());
        self.handlers.push(Box::new(handler))
    }

    /// Start processing messages
    ///
    /// Call process_events on all of the adapters and `recv` on the `IncomingMessage` channel.
    /// Distribute IncomingMessages to list of handlers.
    pub fn run(&self) {

        assert!(self.adapters.len() > 0);
        assert!(self.handlers.len() > 0);

        println!("Chatbot: {} adapters", self.adapters.len());
        println!("Chatbot: {} handlers", self.handlers.len());

        let (incoming_tx, incoming_rx) = channel();

        for adapter in &self.adapters {
            adapter.process_events(self, incoming_tx.clone());
        }

        loop {
            // Get message from adapter
            let msg = match incoming_rx.recv() {
                Ok(msg) => msg,
                Err(_) => break
            };

            // Distribute to handlers
            for handler in &self.handlers {
                let msg_str = msg.get_contents();

                if handler.can_handle(msg_str) {
                    match handler.handle(&msg) {
                        Err(e) => {
                            println!("Error in handler `{}`", handler.name());
                            println!("{:?}", e);
                            println!("The incoming message was {}", msg_str);

                            // TODO remove handler?
                        },
                        _ => ()
                    }
                }
            }
        }

        println!("chatbot shutting down");
    }
}

#[cfg(test)]
mod tests {
    use chatbot::Chatbot;
    use adapter::CliAdapter;

    #[test]
    fn test_create_chatbot() {
        let name = "testbot";
        let bot = Chatbot::new("testbot");
        assert_eq!(bot.get_name(), name);
    }

    #[test]
    fn test_chatbot_add_adapter() {
        let mut bot = Chatbot::new("testbot");
        let cli = CliAdapter::new();
        bot.add_adapter(cli);
    }

    #[test]
    fn test_chatbot_add_handler() {
        let mut bot = Chatbot::new("testbot");
        let echo = handler!("EchoHandler", r"echo .+", |_, msg| {
            Some(msg.to_owned())
        });
        bot.add_handler(echo);
    }
}
<commit_msg>Add filter for when addresser regex is present<commit_after>use regex::Regex;
use std::sync::mpsc::channel;

use adapter::ChatAdapter;
use handler::MessageHandler;

/// The Chatbot is the central data structure of the chatbot platform. It contains a `run` method
/// which listens for messages from adapters and routes them to handlers. Any program which uses
/// chatbot will need to minimally create a Chatbot, add an adapter, add a handler, and call Chatbot
/// [`run`](chatbot/struct.Chatbot.html#method.run).
pub struct Chatbot {
    name: String,
    adapters: Vec<Box<ChatAdapter>>,
    handlers: Vec<Box<MessageHandler>>,
    addresser: Option<Regex>
}

impl Chatbot {
    /// Create a new chatbot instance.
    ///
    /// The name provided here will be used for message filtering and in handlers for whatever they
    /// want. You'll want to make your binding mutable so you can call `add_adapter` and
    /// `add_handler`.
    pub fn new(name: &str) -> Chatbot {
        Chatbot {
            name: name.to_owned(),
            adapters: Vec::new(),
            handlers: Vec::new(),
            addresser: None
        }
    }

    /// Return the name provided in initialization
    pub fn get_name(&self) -> &str {
        self.name.as_ref()
    }

    pub fn get_addresser(&self) -> Option<&Regex> {
        self.addresser.as_ref()
    }

    pub fn address_by_name(&mut self) {
        let addresser_str = format!(r"^\s*@?{}[:,\s]\s*", self.get_name());
        let addresser = Regex::new(addresser_str.as_ref());

        self.addresser = addresser.ok();
    }

    /// Add a ChatAdapter to the bot
    ///
    /// Add as many adapters as you like. The IncomingMessages sent by adapters are made available
    /// to all handlers regardless of how many adapters exist. The IncomingMessage.reply method
    /// makes sure the response is sent back to the adapter from whence the message came.
    pub fn add_adapter<T>(&mut self, adapter: T)
        where T: ChatAdapter + 'static
    {
        println!("Adding adapter {}", adapter.get_name());
        self.adapters.push(Box::new(adapter))
    }

    /// Add a MessageHandler to the bot
    ///
    /// The more handlers you have the more useful your bot becomes (for potentially loose
    /// definitions of useful :P). Check out the handler! macro for making simple handlers and also
    /// see the `MessageHandler` implementors section for a list of built in handlers.
    pub fn add_handler<T>(&mut self, handler: T)
        where T: MessageHandler + 'static
    {
        println!("Adding handler {}", handler.name());
        self.handlers.push(Box::new(handler))
    }

    /// Start processing messages
    ///
    /// Call process_events on all of the adapters and `recv` on the `IncomingMessage` channel.
    /// Distribute IncomingMessages to list of handlers.
    pub fn run(&self) {

        assert!(self.adapters.len() > 0);
        assert!(self.handlers.len() > 0);

        println!("Chatbot: {} adapters", self.adapters.len());
        println!("Chatbot: {} handlers", self.handlers.len());

        let (incoming_tx, incoming_rx) = channel();

        for adapter in &self.adapters {
            adapter.process_events(self, incoming_tx.clone());
        }

        loop {
            // Get message from adapter
            let msg = match incoming_rx.recv() {
                Ok(msg) => msg,
                Err(_) => break
            };

            // Distribute to handlers
            for handler in &self.handlers {
                let msg_str = msg.get_contents();

                if handler.can_handle(msg_str) {
                    if let Some(ref addresser) = self.addresser {
                      if !addresser.is_match(msg_str) {
                        continue;
                      }
                    }

                    match handler.handle(&msg) {
                        Err(e) => {
                            println!("Error in handler `{}`", handler.name());
                            println!("{:?}", e);
                            println!("The incoming message was {}", msg_str);

                            // TODO remove handler?
                        },
                        _ => ()
                    }
                }
            }
        }

        println!("chatbot shutting down");
    }
}

#[cfg(test)]
mod tests {
    use chatbot::Chatbot;
    use adapter::CliAdapter;

    #[test]
    fn test_create_chatbot() {
        let name = "testbot";
        let bot = Chatbot::new("testbot");
        assert_eq!(bot.get_name(), name);
    }

    #[test]
    fn test_chatbot_add_adapter() {
        let mut bot = Chatbot::new("testbot");
        let cli = CliAdapter::new();
        bot.add_adapter(cli);
    }

    #[test]
    fn test_chatbot_add_handler() {
        let mut bot = Chatbot::new("testbot");
        let echo = handler!("EchoHandler", r"echo .+", |_, msg| {
            Some(msg.to_owned())
        });
        bot.add_handler(echo);
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with the chunked format.
//!
//! The chunked file format is extremely simple serialization procedure where
//! each chunk of raw bytes is prepended with the size of the chunk (as 64-bit
//! unsigned big-endian integer).
//!
//! Its primary application is streaming encoding and decoding of blobs of data
//! in formats that do not support or are inefficient for these purposes (such
//! as serialized Protocol Buffer messages).

use std::io::Cursor;

use byteorder::BigEndian;

/// Encodes a given iterator over binary blobs into the chunked format.
///
/// This is a streaming encoder and performs the encoding in a lazy way. It
/// should compose well with other streaming encoders (e.g. these offered by
/// the [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let data = [b"foo", b"bar", b"baz"];
///
/// let mut stream = rrg::chunked::encode(data.iter().map(|blob| &blob[..]));
/// let mut file = File::create("output.chunked").unwrap();
/// std::io::copy(&mut stream, &mut file).unwrap();
/// ```
pub fn encode<I>(iter: I) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    Encode {
        iter: iter,
        cur: Cursor::new(vec!()),
    }
}

/// Decodes a buffer in the chunked format into binary blobs.
///
/// This is a streaming decoder and performs the decoding in a lazy way. It
/// should compose well with other streaming decoders (e.g. these offered by the
/// [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let file = File::open("input.chunked").unwrap();
/// for (idx, blob) in rrg::chunked::decode(file).enumerate() {
///     println!("blob #{}: {:?}", idx, blob.unwrap());
/// }
/// ```
pub fn decode<R, M>(buf: R) -> Decode<R, M>
where
    R: std::io::Read,
    M: prost::Message,
{
    Decode {
        reader: buf,
        buf: vec!(),
        marker: std::marker::PhantomData,
    }
}

/// Streaming encoder for the chunked format.
///
/// It implements the `Read` trait, lazily polling the underlying chunk iterator
/// as more bytes is needed.
///
/// Instances of this type can be constructed using the [`encode`] function.
///
/// [`encode`]: fn.encode.html
pub struct Encode<I> {
    iter: I,
    cur: Cursor<Vec<u8>>,
}

impl<I, M> Encode<I>
where
    I: Iterator<Item = M>,
    M: prost::Message,
{
    /// Checks whether all the data from the underlying cursor has been read.
    fn is_empty(&self) -> bool {
        self.cur.position() == self.cur.get_ref().len() as u64
    }

    /// Pulls another blob of data from the underlying iterator.
    fn pull(&mut self) -> std::io::Result<()> {
        use byteorder::WriteBytesExt as _;

        let msg = match self.iter.next() {
            Some(msg) => msg,
            None => return Ok(()),
        };

        self.cur.get_mut().clear();
        self.cur.set_position(0);

        self.cur.write_u64::<BigEndian>(msg.encoded_len() as u64)?;
        msg.encode(&mut self.cur.get_mut())?;
        self.cur.set_position(0);

        Ok(())
    }
}

impl<I, M> std::io::Read for Encode<I>
where
    I: Iterator<Item = M>,
    M: prost::Message,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.is_empty() {
            self.pull()?;
        }

        self.cur.read(buf)
    }
}

/// Streaming decoder for the chunked format.
///
/// It implements the `Iterator` trait yielding chunks of decoded blobs, lazily
/// decoding data from the underlying buffer.
///
/// Instances of this type can be constructed using the [`decode`] function.
///
/// [`decode`]: fn.decode.html
pub struct Decode<R, M> {
    reader: R,
    buf: Vec<u8>,
    marker: std::marker::PhantomData<M>,
}

impl<R: std::io::Read, M> Decode<R, M> {

    /// Reads a size tag from the underlying buffer.
    ///
    /// It will return `None` if the is no more data in the buffer.
    fn read_len(&mut self) -> std::io::Result<Option<usize>> {
        use byteorder::ReadBytesExt as _;

        // `read` might not always read all 8 bytes. On the other hand, we also
        // cannot use just `read_exact` because the stream might have ended
        // already. Hence, we combine the two. First we attempt to read some
        // bytes with `read`: it should either return 0 (indicating end of the
        // stream), 8 (indicating that we have filled the whole buffer fully)
        // or something in between. In the last case, we use `read_exact to get
        // the remaining bytes (which should be non-zero now).
        let mut buf = [0; 8];
        match self.reader.read(&mut buf[..])? {
            8 => (),
            0 => return Ok(None),
            len => self.reader.read_exact(&mut buf[len..])?,
        }

        let len = (&buf[..]).read_u64::<BigEndian>()? as usize;
        Ok(Some(len))
    }
}

impl<R, M> Iterator for Decode<R, M>
where
    R: std::io::Read,
    M: prost::Message + Default,
{
    type Item = std::io::Result<M>;

    fn next(&mut self) -> Option<std::io::Result<M>> {
        let len = match self.read_len() {
            Ok(Some(len)) => len,
            Ok(None) => return None,
            Err(error) => return Some(Err(error)),
        };

        self.buf.resize(len, u8::default());
        match self.reader.read_exact(&mut self.buf[..]) {
            Ok(()) => (),
            Err(error) => return Some(Err(error)),
        }

        let msg = match M::decode(&self.buf[..]) {
            Ok(msg) => msg,
            Err(error) => return Some(Err(error.into())),
        };

        Some(Ok(msg))
    }
}

#[cfg(test)]
pub mod tests {

    use super::*;

    #[test]
    pub fn test_encode_empty_iter() {
        use std::io::Read as _;

        let mut stream = encode(std::iter::empty::<()>());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert!(output.is_empty());
    }

    #[test]
    pub fn test_decode_empty_buf() {
        let buf: &[u8] = b"";
        let mut iter = decode::<_, ()>(buf);

        assert!(iter.next().is_none());
    }

    #[test]
    pub fn test_decode_incorrect_size_tag() {
        let buf: &[u8] = b"\x12\x34\x56";
        let mut iter = decode::<_, ()>(buf);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_decode_zero_size_tag() {
        let buf: &[u8] = b"\x00\x00\x00\x00\x00\x00\x00\x00";

        let mut iter = decode(buf).map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_partial_size_tag() {

        // A simple reader that yields a 0-valued size tag byte by byte.
        struct Reader(u8);

        impl std::io::Read for Reader {

            fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
                dbg!(self.0);
                if self.0 == 8 {
                    Ok(0)
                } else {
                    buf[0] = 0;
                    self.0 += 1;
                    Ok(1)
                }
            }
        }

        let mut iter = decode(Reader(0)).map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_short_input() {
        let buf: &[u8] = b"\x00\x00\x00\x00\x00\x00\x00\x42foo";
        let mut iter = decode::<_, ()>(buf);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_encode_and_decode_single_message() {
        let mut iter = decode(encode(std::iter::once(String::from("foo"))))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_single_unit_message() {
        let mut iter = decode(encode(std::iter::once(())))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_messages() {
        let msgs = vec! {
            b"foo".to_vec(),
            b"bar".to_vec(),
            b"baz".to_vec(),
        };

        let mut iter = decode(encode(msgs.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"foo".to_vec()));
        assert_eq!(iter.next(), Some(b"bar".to_vec()));
        assert_eq!(iter.next(), Some(b"baz".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_unit_messages() {
        let msgs = vec!((), (), ());

        let mut iter = decode(encode(msgs.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }
}
<commit_msg>Remove a spurious debugging leftover.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with the chunked format.
//!
//! The chunked file format is extremely simple serialization procedure where
//! each chunk of raw bytes is prepended with the size of the chunk (as 64-bit
//! unsigned big-endian integer).
//!
//! Its primary application is streaming encoding and decoding of blobs of data
//! in formats that do not support or are inefficient for these purposes (such
//! as serialized Protocol Buffer messages).

use std::io::Cursor;

use byteorder::BigEndian;

/// Encodes a given iterator over binary blobs into the chunked format.
///
/// This is a streaming encoder and performs the encoding in a lazy way. It
/// should compose well with other streaming encoders (e.g. these offered by
/// the [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let data = [b"foo", b"bar", b"baz"];
///
/// let mut stream = rrg::chunked::encode(data.iter().map(|blob| &blob[..]));
/// let mut file = File::create("output.chunked").unwrap();
/// std::io::copy(&mut stream, &mut file).unwrap();
/// ```
pub fn encode<I>(iter: I) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    Encode {
        iter: iter,
        cur: Cursor::new(vec!()),
    }
}

/// Decodes a buffer in the chunked format into binary blobs.
///
/// This is a streaming decoder and performs the decoding in a lazy way. It
/// should compose well with other streaming decoders (e.g. these offered by the
/// [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let file = File::open("input.chunked").unwrap();
/// for (idx, blob) in rrg::chunked::decode(file).enumerate() {
///     println!("blob #{}: {:?}", idx, blob.unwrap());
/// }
/// ```
pub fn decode<R, M>(buf: R) -> Decode<R, M>
where
    R: std::io::Read,
    M: prost::Message,
{
    Decode {
        reader: buf,
        buf: vec!(),
        marker: std::marker::PhantomData,
    }
}

/// Streaming encoder for the chunked format.
///
/// It implements the `Read` trait, lazily polling the underlying chunk iterator
/// as more bytes is needed.
///
/// Instances of this type can be constructed using the [`encode`] function.
///
/// [`encode`]: fn.encode.html
pub struct Encode<I> {
    iter: I,
    cur: Cursor<Vec<u8>>,
}

impl<I, M> Encode<I>
where
    I: Iterator<Item = M>,
    M: prost::Message,
{
    /// Checks whether all the data from the underlying cursor has been read.
    fn is_empty(&self) -> bool {
        self.cur.position() == self.cur.get_ref().len() as u64
    }

    /// Pulls another blob of data from the underlying iterator.
    fn pull(&mut self) -> std::io::Result<()> {
        use byteorder::WriteBytesExt as _;

        let msg = match self.iter.next() {
            Some(msg) => msg,
            None => return Ok(()),
        };

        self.cur.get_mut().clear();
        self.cur.set_position(0);

        self.cur.write_u64::<BigEndian>(msg.encoded_len() as u64)?;
        msg.encode(&mut self.cur.get_mut())?;
        self.cur.set_position(0);

        Ok(())
    }
}

impl<I, M> std::io::Read for Encode<I>
where
    I: Iterator<Item = M>,
    M: prost::Message,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.is_empty() {
            self.pull()?;
        }

        self.cur.read(buf)
    }
}

/// Streaming decoder for the chunked format.
///
/// It implements the `Iterator` trait yielding chunks of decoded blobs, lazily
/// decoding data from the underlying buffer.
///
/// Instances of this type can be constructed using the [`decode`] function.
///
/// [`decode`]: fn.decode.html
pub struct Decode<R, M> {
    reader: R,
    buf: Vec<u8>,
    marker: std::marker::PhantomData<M>,
}

impl<R: std::io::Read, M> Decode<R, M> {

    /// Reads a size tag from the underlying buffer.
    ///
    /// It will return `None` if the is no more data in the buffer.
    fn read_len(&mut self) -> std::io::Result<Option<usize>> {
        use byteorder::ReadBytesExt as _;

        // `read` might not always read all 8 bytes. On the other hand, we also
        // cannot use just `read_exact` because the stream might have ended
        // already. Hence, we combine the two. First we attempt to read some
        // bytes with `read`: it should either return 0 (indicating end of the
        // stream), 8 (indicating that we have filled the whole buffer fully)
        // or something in between. In the last case, we use `read_exact to get
        // the remaining bytes (which should be non-zero now).
        let mut buf = [0; 8];
        match self.reader.read(&mut buf[..])? {
            8 => (),
            0 => return Ok(None),
            len => self.reader.read_exact(&mut buf[len..])?,
        }

        let len = (&buf[..]).read_u64::<BigEndian>()? as usize;
        Ok(Some(len))
    }
}

impl<R, M> Iterator for Decode<R, M>
where
    R: std::io::Read,
    M: prost::Message + Default,
{
    type Item = std::io::Result<M>;

    fn next(&mut self) -> Option<std::io::Result<M>> {
        let len = match self.read_len() {
            Ok(Some(len)) => len,
            Ok(None) => return None,
            Err(error) => return Some(Err(error)),
        };

        self.buf.resize(len, u8::default());
        match self.reader.read_exact(&mut self.buf[..]) {
            Ok(()) => (),
            Err(error) => return Some(Err(error)),
        }

        let msg = match M::decode(&self.buf[..]) {
            Ok(msg) => msg,
            Err(error) => return Some(Err(error.into())),
        };

        Some(Ok(msg))
    }
}

#[cfg(test)]
pub mod tests {

    use super::*;

    #[test]
    pub fn test_encode_empty_iter() {
        use std::io::Read as _;

        let mut stream = encode(std::iter::empty::<()>());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert!(output.is_empty());
    }

    #[test]
    pub fn test_decode_empty_buf() {
        let buf: &[u8] = b"";
        let mut iter = decode::<_, ()>(buf);

        assert!(iter.next().is_none());
    }

    #[test]
    pub fn test_decode_incorrect_size_tag() {
        let buf: &[u8] = b"\x12\x34\x56";
        let mut iter = decode::<_, ()>(buf);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_decode_zero_size_tag() {
        let buf: &[u8] = b"\x00\x00\x00\x00\x00\x00\x00\x00";

        let mut iter = decode(buf).map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_partial_size_tag() {

        // A simple reader that yields a 0-valued size tag byte by byte.
        struct Reader(u8);

        impl std::io::Read for Reader {

            fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
                if self.0 == 8 {
                    Ok(0)
                } else {
                    buf[0] = 0;
                    self.0 += 1;
                    Ok(1)
                }
            }
        }

        let mut iter = decode(Reader(0)).map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_short_input() {
        let buf: &[u8] = b"\x00\x00\x00\x00\x00\x00\x00\x42foo";
        let mut iter = decode::<_, ()>(buf);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_encode_and_decode_single_message() {
        let mut iter = decode(encode(std::iter::once(String::from("foo"))))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_single_unit_message() {
        let mut iter = decode(encode(std::iter::once(())))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_messages() {
        let msgs = vec! {
            b"foo".to_vec(),
            b"bar".to_vec(),
            b"baz".to_vec(),
        };

        let mut iter = decode(encode(msgs.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"foo".to_vec()));
        assert_eq!(iter.next(), Some(b"bar".to_vec()));
        assert_eq!(iter.next(), Some(b"baz".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_unit_messages() {
        let msgs = vec!((), (), ());

        let mut iter = decode(encode(msgs.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), Some(()));
        assert_eq!(iter.next(), None);
    }
}
<|endoftext|>


<commit_before>extern crate encoding;

use encoding::Encoding;
use encoding::DecoderTrap;
use encoding::all::ISO_8859_1;
use encoding::types::RawDecoder;
use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::fmt::Display;
use std::fmt::Formatter;
use std::io;
use std::io::BufRead;
use std::io::Bytes;
use std::io::ErrorKind;
use std::io::Lines;
use std::io::Read;
use std::iter::Peekable;

/////////////////////

#[derive(Debug)]
pub struct PropertiesError {
  description: String,
  cause: Option<Box<Error+Send+Sync>>,
}

impl PropertiesError {
  fn new(description: &str, cause: Option<Box<Error+Send+Sync>>) -> Self {
    PropertiesError {
      description: description.to_string(),
      cause: cause,
    }
  }
}

impl Error for PropertiesError {
  fn description(&self) -> &str {
    &self.description
  }
}

impl Display for PropertiesError {
  fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    Display::fmt(&self.description, f)
  }
}

/////////////////////

#[derive(PartialEq,Eq,Debug)]
struct NaturalLine(String);

// We can't use BufRead.lines() because it doesn't use the proper line endings
struct NaturalLines<R: Read> {
  bytes: Peekable<Bytes<R>>,
  eof: bool,
}

impl<R: Read> NaturalLines<R> {
  fn new(reader: R) -> Self {
    NaturalLines {
      bytes: reader.bytes().peekable(),
      eof: false,
    }
  }
}

const LF: u8 = 10;
const CR: u8 = 13;

impl< R: Read> Iterator for NaturalLines<R> {
  type Item = io::Result<NaturalLine>;

  fn next(&mut self) -> Option<Self::Item> {
    if self.eof {
      return None;
    }
    let mut buf = String::new();
    loop {
      match self.bytes.next() {
        Some(r) => match r {
          Ok(b) => {
            match b {
              CR => {
                match self.bytes.peek() {
                  Some(&Ok(LF)) => { self.bytes.next(); },
                  _ => (),
                }
                return Some(Ok(NaturalLine(buf)));
              },
              LF => return Some(Ok(NaturalLine(buf))),
              _ => match ISO_8859_1.decode(&[b], DecoderTrap::Strict) {
                Ok(s) => buf.push_str(&s),
                Err(e) => return Some(Err(io::Error::new(ErrorKind::InvalidData, "Error reading ISO-8859-1 encoding"))),
              },
            };
          },
          Err(e) => return Some(Err(e)),
        },
        None => {
          self.eof = true;
          return Some(Ok(NaturalLine(buf)));
        },
      }
    }
  }
}

/////////////////////

#[derive(PartialEq,Eq,Debug)]
struct LogicalLine(String);

// We can't use BufRead.lines() because it doesn't use the proper line endings
struct LogicalLines<I: Iterator<Item=io::Result<NaturalLine>>> {
  physical_lines: I,
  eof: bool,
}

impl<I: Iterator<Item=io::Result<NaturalLine>>> LogicalLines<I> {
  fn new(physical_lines: I) -> Self {
    LogicalLines {
      physical_lines: physical_lines,
      eof: false,
    }
  }
}

fn count_ending_backslashes(s: &str) -> usize {
  let chars: Vec<char> = s.chars().collect();
  let n = chars.len();
  let mut i = n;
  while i > 0 {
    i -= 1;
    if chars[i] != '\\' {
      return n - 1 - i;
    }
  }
  n
}

impl<I: Iterator<Item=io::Result<NaturalLine>>> Iterator for LogicalLines<I> {
  type Item = io::Result<LogicalLine>;

  fn next(&mut self) -> Option<Self::Item> {
    if self.eof {
      return None;
    }
    let mut buf = String::new();
    loop {
      match self.physical_lines.next() {
        Some(Err(e)) => return Some(Err(e)),
        Some(Ok(NaturalLine(mut line))) => {
          buf.push_str(&line);
          if count_ending_backslashes(&line) % 2 == 1 {
            buf.pop();
          } else {
            return Some(Ok(LogicalLine(buf)));
          }
        },
        None => {
          self.eof = true;
          return None;
        },
      }
    }
  }
}

/////////////////////

pub fn load(reader: &mut BufRead) -> Result<HashMap<String, String>, PropertiesError> {
  let map = HashMap::new();
  let mut lines = LogicalLines::new(NaturalLines::new(reader));
  loop {
    match lines.next() {
      Some(line) => (), // TODO
      None => break,
    }
  }
  Ok(map)
}

/////////////////////

#[cfg(test)]
mod tests {
  use std::io;
  use std::io::BufRead;
  use std::io::BufReader;
  use std::io::Read;
  use super::CR;
  use super::LF;
  use super::LogicalLine;
  use super::LogicalLines;
  use super::NaturalLine;
  use super::NaturalLines;

  const BS: u8 = 92; // backslash
  const SP: u8 = 32; // space

  #[test]
  fn natural_lines() {
    let data = [
      (vec![], vec![""]),
      (vec![SP], vec![" "]),
      (vec![SP, CR], vec![" ", ""]),
      (vec![SP, LF], vec![" ", ""]),
      (vec![SP, CR, LF], vec![" ", ""]),
      (vec![SP, CR, SP], vec![" ", " "]),
      (vec![SP, LF, SP], vec![" ", " "]),
      (vec![SP, CR, LF, SP], vec![" ", " "]),
      (vec![CR], vec!["", ""]),
      (vec![LF], vec!["", ""]),
      (vec![CR, LF], vec!["", ""]),
      (vec![CR, SP], vec!["", " "]),
      (vec![LF, SP], vec!["", " "]),
      (vec![CR, LF, SP], vec!["", " "]),
    ];
    for &(ref bytes, ref lines) in data.iter() {
      let reader = &bytes as &[u8];
      let mut iter = NaturalLines::new(reader);
      for line in lines {
        match (line.to_string(), iter.next()) {
          (e, Some(Ok(NaturalLine(a)))) => if e != a {
            panic!("Failure while processing {:?}.  Expected Some(Ok({:?})), but was {:?}", bytes, e, a);
          },
          (e, a) => panic!("Failure while processing {:?}.  Expected Some(Ok({:?})), but was {:?}", bytes, e, a),
        }
      }
      match iter.next() {
        None => (),
        a => panic!("Failure while processing {:?}.  Expected None, but was {:?}", bytes, a),
      }
    }
  }

  #[test]
  fn logical_lines() {
    let data = [
      (vec![], vec![]),
      (vec!["foo"], vec!["foo"]),
      (vec!["foo", "bar"], vec!["foo", "bar"]),
      (vec!["foo\\", "bar"], vec!["foobar"]),
      (vec!["foo\\\\", "bar"], vec!["foo\\\\", "bar"]),
      (vec!["foo\\\\\\", "bar"], vec!["foo\\\\bar"]),
    ];
    for &(ref input_lines, ref lines) in data.iter() {
      let input_natural_lines: Vec<io::Result<NaturalLine>> =
            input_lines.iter().map(|x| Ok(NaturalLine(x.to_string()))).collect();
      let mut iter = LogicalLines::new(input_lines.iter().map(|x| Ok(NaturalLine(x.to_string()))));
      for line in lines {
        match (line.to_string(), iter.next()) {
          (e, Some(Ok(LogicalLine(a)))) => if e != a {
            panic!("Failure while processing {:?}.  Expected Some(Ok({:?})), but was {:?}", input_lines, e, a);
          },
          (e, a) => panic!("Failure while processing {:?}.  Expected Some(Ok({:?})), but was {:?}", input_lines, e, a),
        }
      }
      match iter.next() {
        None => (),
        a => panic!("Failure while processing {:?}.  Expected None, but was {:?}", input_lines, a),
      }
    }
  }

  #[test]
  fn count_ending_backslashes() {
    assert_eq!(0, super::count_ending_backslashes(""));

    assert_eq!(0, super::count_ending_backslashes("x"));
    assert_eq!(1, super::count_ending_backslashes("\\"));

    assert_eq!(0, super::count_ending_backslashes("xx"));
    assert_eq!(0, super::count_ending_backslashes("\\x"));
    assert_eq!(1, super::count_ending_backslashes("x\\"));
    assert_eq!(2, super::count_ending_backslashes("\\\\"));

    assert_eq!(0, super::count_ending_backslashes("xxx"));
    assert_eq!(0, super::count_ending_backslashes("\\xx"));
    assert_eq!(0, super::count_ending_backslashes("x\\x"));
    assert_eq!(0, super::count_ending_backslashes("\\\\x"));
    assert_eq!(1, super::count_ending_backslashes("xx\\"));
    assert_eq!(1, super::count_ending_backslashes("\\x\\"));
    assert_eq!(2, super::count_ending_backslashes("x\\\\"));
    assert_eq!(3, super::count_ending_backslashes("\\\\\\"));
  }
}
<commit_msg>Remove extraneous comment<commit_after>extern crate encoding;

use encoding::Encoding;
use encoding::DecoderTrap;
use encoding::all::ISO_8859_1;
use encoding::types::RawDecoder;
use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::fmt::Display;
use std::fmt::Formatter;
use std::io;
use std::io::BufRead;
use std::io::Bytes;
use std::io::ErrorKind;
use std::io::Lines;
use std::io::Read;
use std::iter::Peekable;

/////////////////////

#[derive(Debug)]
pub struct PropertiesError {
  description: String,
  cause: Option<Box<Error+Send+Sync>>,
}

impl PropertiesError {
  fn new(description: &str, cause: Option<Box<Error+Send+Sync>>) -> Self {
    PropertiesError {
      description: description.to_string(),
      cause: cause,
    }
  }
}

impl Error for PropertiesError {
  fn description(&self) -> &str {
    &self.description
  }
}

impl Display for PropertiesError {
  fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    Display::fmt(&self.description, f)
  }
}

/////////////////////

#[derive(PartialEq,Eq,Debug)]
struct NaturalLine(String);

// We can't use BufRead.lines() because it doesn't use the proper line endings
struct NaturalLines<R: Read> {
  bytes: Peekable<Bytes<R>>,
  eof: bool,
}

impl<R: Read> NaturalLines<R> {
  fn new(reader: R) -> Self {
    NaturalLines {
      bytes: reader.bytes().peekable(),
      eof: false,
    }
  }
}

const LF: u8 = 10;
const CR: u8 = 13;

impl< R: Read> Iterator for NaturalLines<R> {
  type Item = io::Result<NaturalLine>;

  fn next(&mut self) -> Option<Self::Item> {
    if self.eof {
      return None;
    }
    let mut buf = String::new();
    loop {
      match self.bytes.next() {
        Some(r) => match r {
          Ok(b) => {
            match b {
              CR => {
                match self.bytes.peek() {
                  Some(&Ok(LF)) => { self.bytes.next(); },
                  _ => (),
                }
                return Some(Ok(NaturalLine(buf)));
              },
              LF => return Some(Ok(NaturalLine(buf))),
              _ => match ISO_8859_1.decode(&[b], DecoderTrap::Strict) {
                Ok(s) => buf.push_str(&s),
                Err(e) => return Some(Err(io::Error::new(ErrorKind::InvalidData, "Error reading ISO-8859-1 encoding"))),
              },
            };
          },
          Err(e) => return Some(Err(e)),
        },
        None => {
          self.eof = true;
          return Some(Ok(NaturalLine(buf)));
        },
      }
    }
  }
}

/////////////////////

#[derive(PartialEq,Eq,Debug)]
struct LogicalLine(String);

struct LogicalLines<I: Iterator<Item=io::Result<NaturalLine>>> {
  physical_lines: I,
  eof: bool,
}

impl<I: Iterator<Item=io::Result<NaturalLine>>> LogicalLines<I> {
  fn new(physical_lines: I) -> Self {
    LogicalLines {
      physical_lines: physical_lines,
      eof: false,
    }
  }
}

fn count_ending_backslashes(s: &str) -> usize {
  let chars: Vec<char> = s.chars().collect();
  let n = chars.len();
  let mut i = n;
  while i > 0 {
    i -= 1;
    if chars[i] != '\\' {
      return n - 1 - i;
    }
  }
  n
}

impl<I: Iterator<Item=io::Result<NaturalLine>>> Iterator for LogicalLines<I> {
  type Item = io::Result<LogicalLine>;

  fn next(&mut self) -> Option<Self::Item> {
    if self.eof {
      return None;
    }
    let mut buf = String::new();
    loop {
      match self.physical_lines.next() {
        Some(Err(e)) => return Some(Err(e)),
        Some(Ok(NaturalLine(mut line))) => {
          buf.push_str(&line);
          if count_ending_backslashes(&line) % 2 == 1 {
            buf.pop();
          } else {
            return Some(Ok(LogicalLine(buf)));
          }
        },
        None => {
          self.eof = true;
          return None;
        },
      }
    }
  }
}

/////////////////////

pub fn load(reader: &mut BufRead) -> Result<HashMap<String, String>, PropertiesError> {
  let map = HashMap::new();
  let mut lines = LogicalLines::new(NaturalLines::new(reader));
  loop {
    match lines.next() {
      Some(line) => (), // TODO
      None => break,
    }
  }
  Ok(map)
}

/////////////////////

#[cfg(test)]
mod tests {
  use std::io;
  use std::io::BufRead;
  use std::io::BufReader;
  use std::io::Read;
  use super::CR;
  use super::LF;
  use super::LogicalLine;
  use super::LogicalLines;
  use super::NaturalLine;
  use super::NaturalLines;

  const BS: u8 = 92; // backslash
  const SP: u8 = 32; // space

  #[test]
  fn natural_lines() {
    let data = [
      (vec![], vec![""]),
      (vec![SP], vec![" "]),
      (vec![SP, CR], vec![" ", ""]),
      (vec![SP, LF], vec![" ", ""]),
      (vec![SP, CR, LF], vec![" ", ""]),
      (vec![SP, CR, SP], vec![" ", " "]),
      (vec![SP, LF, SP], vec![" ", " "]),
      (vec![SP, CR, LF, SP], vec![" ", " "]),
      (vec![CR], vec!["", ""]),
      (vec![LF], vec!["", ""]),
      (vec![CR, LF], vec!["", ""]),
      (vec![CR, SP], vec!["", " "]),
      (vec![LF, SP], vec!["", " "]),
      (vec![CR, LF, SP], vec!["", " "]),
    ];
    for &(ref bytes, ref lines) in data.iter() {
      let reader = &bytes as &[u8];
      let mut iter = NaturalLines::new(reader);
      for line in lines {
        match (line.to_string(), iter.next()) {
          (e, Some(Ok(NaturalLine(a)))) => if e != a {
            panic!("Failure while processing {:?}.  Expected Some(Ok({:?})), but was {:?}", bytes, e, a);
          },
          (e, a) => panic!("Failure while processing {:?}.  Expected Some(Ok({:?})), but was {:?}", bytes, e, a),
        }
      }
      match iter.next() {
        None => (),
        a => panic!("Failure while processing {:?}.  Expected None, but was {:?}", bytes, a),
      }
    }
  }

  #[test]
  fn logical_lines() {
    let data = [
      (vec![], vec![]),
      (vec!["foo"], vec!["foo"]),
      (vec!["foo", "bar"], vec!["foo", "bar"]),
      (vec!["foo\\", "bar"], vec!["foobar"]),
      (vec!["foo\\\\", "bar"], vec!["foo\\\\", "bar"]),
      (vec!["foo\\\\\\", "bar"], vec!["foo\\\\bar"]),
    ];
    for &(ref input_lines, ref lines) in data.iter() {
      let input_natural_lines: Vec<io::Result<NaturalLine>> =
            input_lines.iter().map(|x| Ok(NaturalLine(x.to_string()))).collect();
      let mut iter = LogicalLines::new(input_lines.iter().map(|x| Ok(NaturalLine(x.to_string()))));
      for line in lines {
        match (line.to_string(), iter.next()) {
          (e, Some(Ok(LogicalLine(a)))) => if e != a {
            panic!("Failure while processing {:?}.  Expected Some(Ok({:?})), but was {:?}", input_lines, e, a);
          },
          (e, a) => panic!("Failure while processing {:?}.  Expected Some(Ok({:?})), but was {:?}", input_lines, e, a),
        }
      }
      match iter.next() {
        None => (),
        a => panic!("Failure while processing {:?}.  Expected None, but was {:?}", input_lines, a),
      }
    }
  }

  #[test]
  fn count_ending_backslashes() {
    assert_eq!(0, super::count_ending_backslashes(""));

    assert_eq!(0, super::count_ending_backslashes("x"));
    assert_eq!(1, super::count_ending_backslashes("\\"));

    assert_eq!(0, super::count_ending_backslashes("xx"));
    assert_eq!(0, super::count_ending_backslashes("\\x"));
    assert_eq!(1, super::count_ending_backslashes("x\\"));
    assert_eq!(2, super::count_ending_backslashes("\\\\"));

    assert_eq!(0, super::count_ending_backslashes("xxx"));
    assert_eq!(0, super::count_ending_backslashes("\\xx"));
    assert_eq!(0, super::count_ending_backslashes("x\\x"));
    assert_eq!(0, super::count_ending_backslashes("\\\\x"));
    assert_eq!(1, super::count_ending_backslashes("xx\\"));
    assert_eq!(1, super::count_ending_backslashes("\\x\\"));
    assert_eq!(2, super::count_ending_backslashes("x\\\\"));
    assert_eq!(3, super::count_ending_backslashes("\\\\\\"));
  }
}
<|endoftext|>


<commit_before>mod byte_conversion;
mod byte_manipulation;
mod byte_utilities;
mod english_scoring;
mod key;

fn main() {

}<commit_msg>Ignore unused functions during cargo test<commit_after>#![allow(dead_code)]

mod byte_conversion;
mod byte_manipulation;
mod byte_utilities;
mod english_scoring;
mod key;

fn main() {

}<|endoftext|>


<commit_before>#[cfg(test)]
mod test;

use smallvec::SmallVec;

use std::convert::{From, Into};
use std::fmt::{self, Debug, Formatter};
use std::iter::FromIterator;

/// A data-structure for storing a sequence of 4-bit values.
///
/// Values are stored in a `Vec<u8>`, with two values per byte.
///
/// Values at even indices are stored in the most-significant half of their byte,
/// while values at odd indices are stored in the least-significant half.
///
/// Imagine a vector of [MSB][msb-wiki] first bytes, and you'll be right.
///
/// n = [_ _ | _ _ | _ _]
///
/// [msb-wiki]: http://en.wikipedia.org/wiki/Most_significant_bit
#[derive(Clone, Default)]
pub struct NibbleVec {
    length: usize,
    data: SmallVec<[u8; 20]>,
}
// TODO add trait to make size generic
impl NibbleVec {
    /// Create an empty nibble vector.
    pub fn new() -> NibbleVec {
        NibbleVec {
            length: 0,
            data: SmallVec::new(),
        }
    }

    /// Create a nibble vector from a vector of bytes.
    ///
    /// Each byte is split into two 4-bit entries (MSB, LSB).
    #[inline]
    pub fn from_byte_vec(vec: Vec<u8>) -> NibbleVec {
        let length = 2 * vec.len();
        NibbleVec {
            length,
            data: SmallVec::from_iter(vec),
        }
    }

    /// Create a nibble vector from a `SmallVec` of Size.
    ///
    /// Each byte is split into two 4-bit entries (MSB, LSB).
    #[inline]
    pub fn from_small_vec(length: usize) -> NibbleVec {
        NibbleVec {
            length,
            data: SmallVec::new(),
        }
    }

    /// Returns a byte slice of the nibble vector's contents.
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        &self.data[..]
    }

    /// Converts a nibble vector into a byte vector.
    ///
    /// This consumes the nibble vector, so we do not need to copy its contents.
    #[inline]
    pub fn into_bytes(self) -> Vec<u8> {
        self.data.to_vec()
    }

    /// Get the number of elements stored in the vector.
    #[inline]
    pub fn len(&self) -> usize {
        self.length
    }

    /// Returns `true` if the nibble vector has a length of 0.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Fetch a single entry from the vector.
    ///
    /// Guaranteed to be a value in the interval [0, 15].
    ///
    /// **Panics** if `idx >= self.len()`.
    #[inline]
    pub fn get(&self, idx: usize) -> u8 {
        if idx >= self.length {
            panic!(
                "attempted access beyond vector end. len is {}, index is {}",
                self.length, idx
            );
        }
        let vec_idx = idx / 2;
        match idx % 2 {
            // If the index is even, take the first (most significant) half of the stored byte.
            0 => self.data[vec_idx] >> 4,
            // If the index is odd, take the second (least significant) half.
            _ => self.data[vec_idx] & 0x0F,
        }
    }

    /// Add a single nibble to the vector.
    ///
    /// Only the 4 least-significant bits of the value are used.
    #[inline]
    pub fn push(&mut self, val: u8) {
        if self.length % 2 == 0 {
            self.data.push(val << 4);
        } else {
            let vec_len = self.data.len();

            // Zero the second half of the last byte just to be safe.
            self.data[vec_len - 1] &= 0xF0;

            // Write the new value.
            self.data[vec_len - 1] |= val & 0x0F;
        }
        self.length += 1;
    }

    /// Split the vector into two parts.
    ///
    /// All elements at or following the given index are returned in a new `NibbleVec`,
    /// with exactly `idx` elements remaining in this vector.
    ///
    /// **Panics** if `idx > self.len()`.
    pub fn split(&mut self, idx: usize) -> NibbleVec {
        if idx > self.length {
            panic!(
                "attempted to split past vector end. len is {}, index is {}",
                self.length, idx
            );
        } else if idx == self.length {
            NibbleVec::new()
        } else if idx % 2 == 0 {
            self.split_even(idx)
        } else {
            self.split_odd(idx)
        }
    }

    /// Split function for odd *indices*.
    #[inline]
    fn split_odd(&mut self, idx: usize) -> NibbleVec {
        let tail_vec_size = (self.length - idx) / 2;
        let mut tail = NibbleVec::from_small_vec(tail_vec_size);

        // Perform an overlap copy, copying the last nibble of the original vector only if
        // the length of the new tail is *odd*.
        let tail_length = self.length - idx;
        let take_last = tail_length % 2 == 1;
        self.overlap_copy(
            idx / 2,
            self.data.len(),
            &mut tail.data,
            &mut tail.length,
            take_last,
        );

        // Remove the copied bytes, being careful to skip the idx byte.
        for _ in (idx / 2 + 1)..self.data.len() {
            self.data.pop();
        }

        // Zero the second half of the index byte so as to maintain the last-nibble invariant.
        self.data[idx / 2] &= 0xF0;

        // Update the length of the first NibbleVec.
        self.length = idx;

        tail
    }

    /// Split function for even *indices*.
    #[inline]
    fn split_even(&mut self, idx: usize) -> NibbleVec {
        // Avoid allocating a temporary vector by copying all the bytes in order, then popping them.

        // Possible to prove: l_d - ⌊i / 2⌋ = ⌊(l_v - i + 1) / 2⌋
        //  where l_d = self.data.len()
        //        l_v = self.length
        let tail_vec_size = (self.length - idx + 1) / 2;
        let half_idx = idx / 2;
        let mut tail = NibbleVec::from_small_vec(tail_vec_size);

        // Copy the bytes.
        for i in half_idx..self.data.len() {
            tail.data.push(self.data[i]);
        }

        // Pop the same bytes.
        for _ in half_idx..self.data.len() {
            self.data.pop();
        }

        // Update lengths.
        tail.length = self.length - idx;
        self.length = idx;

        tail
    }

    /// Copy data between the second half of self.data[start] and
    /// self.data[end - 1]. The second half of the last entry is included
    /// if include_last is true.
    #[inline]
    fn overlap_copy(
        &self,
        start: usize,
        end: usize,
        vec: &mut SmallVec<[u8; 20]>,
        length: &mut usize,
        include_last: bool,
    ) {
        // Copy up to the first half of the last byte.
        for i in start..(end - 1) {
            // The first half is the second half of the old entry.
            let first_half = self.data[i] & 0x0f;

            // The second half is the first half of the next entry.
            let second_half = self.data[i + 1] >> 4;

            vec.push((first_half << 4) | second_half);
            *length += 2;
        }

        if include_last {
            let last = self.data[end - 1] & 0x0f;
            vec.push(last << 4);
            *length += 1;
        }
    }

    /// Append another nibble vector whilst consuming this vector.
    #[inline]
    pub fn join(mut self, other: &NibbleVec) -> NibbleVec {
        // If the length is even, we can append directly.
        if self.length % 2 == 0 {
            self.length += other.length;
            self.data.extend_from_slice(&other.data);
            return self;
        }

        // If the other vector is empty, bail out.
        if other.is_empty() {
            return self;
        }

        // If the length is odd, we have to perform an overlap copy.
        // Copy the first half of the first element, to make the vector an even length.
        self.push(other.get(0));

        // Copy the rest of the vector using an overlap copy.
        let take_last = other.len() % 2 == 0;
        other.overlap_copy(
            0,
            other.data.len(),
            &mut self.data,
            &mut self.length,
            take_last,
        );

        self
    }
}

impl PartialEq<NibbleVec> for NibbleVec {
    #[inline]
    fn eq(&self, other: &NibbleVec) -> bool {
        self.length == other.length && self.data == other.data
    }
}

impl Eq for NibbleVec {}

/// Compare a `NibbleVec` and a slice of bytes *element-by-element*.
/// Bytes are **not** interpreted as two `NibbleVec` entries.
impl PartialEq<[u8]> for NibbleVec {
    #[inline]
    fn eq(&self, other: &[u8]) -> bool {
        if other.len() != self.len() {
            return false;
        }

        for (i, x) in other.iter().enumerate() {
            if self.get(i) != *x {
                return false;
            }
        }
        true
    }
}

impl Debug for NibbleVec {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "NibbleVec [")?;

        if !self.is_empty() {
            write!(fmt, "{}", self.get(0))?;
        }

        for i in 1..self.len() {
            write!(fmt, ", {}", self.get(i))?;
        }
        write!(fmt, "]")
    }
}

impl From<Vec<u8>> for NibbleVec {
    #[inline]
    fn from(v: Vec<u8>) -> NibbleVec {
        NibbleVec::from_byte_vec(v)
    }
}

impl<'a> From<&'a [u8]> for NibbleVec {
    #[inline]
    fn from(v: &[u8]) -> NibbleVec {
        NibbleVec::from_byte_vec(v.into())
    }
}

impl Into<Vec<u8>> for NibbleVec {
    #[inline]
    fn into(self) -> Vec<u8> {
        self.data.to_vec()
    }
}

impl<'a> Into<Vec<u8>> for &'a NibbleVec {
    #[inline]
    fn into(self) -> Vec<u8> {
        self.data.to_vec()
    }
}
<commit_msg>oops fix split_odd split_even<commit_after>#[cfg(test)]
mod test;

use smallvec::SmallVec;

use std::convert::{From, Into};
use std::fmt::{self, Debug, Formatter};
use std::iter::FromIterator;

/// A data-structure for storing a sequence of 4-bit values.
///
/// Values are stored in a `Vec<u8>`, with two values per byte.
///
/// Values at even indices are stored in the most-significant half of their byte,
/// while values at odd indices are stored in the least-significant half.
///
/// Imagine a vector of [MSB][msb-wiki] first bytes, and you'll be right.
///
/// n = [_ _ | _ _ | _ _]
///
/// [msb-wiki]: http://en.wikipedia.org/wiki/Most_significant_bit
#[derive(Clone, Default)]
pub struct NibbleVec {
    length: usize,
    data: SmallVec<[u8; 20]>,
}
// TODO add trait to make size generic
impl NibbleVec {
    /// Create an empty nibble vector.
    pub fn new() -> NibbleVec {
        NibbleVec {
            length: 0,
            data: SmallVec::new(),
        }
    }

    /// Create a nibble vector from a vector of bytes.
    ///
    /// Each byte is split into two 4-bit entries (MSB, LSB).
    #[inline]
    pub fn from_byte_vec(vec: Vec<u8>) -> NibbleVec {
        let length = 2 * vec.len();
        NibbleVec {
            length,
            data: SmallVec::from_iter(vec),
        }
    }

    /// Create a nibble vector from a `SmallVec` of Size.
    ///
    /// Each byte is split into two 4-bit entries (MSB, LSB).
    #[inline]
    pub fn from_small_vec(length: usize) -> NibbleVec {
        NibbleVec {
            length,
            data: SmallVec::new(),
        }
    }

    /// Returns a byte slice of the nibble vector's contents.
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        &self.data[..]
    }

    /// Converts a nibble vector into a byte vector.
    ///
    /// This consumes the nibble vector, so we do not need to copy its contents.
    #[inline]
    pub fn into_bytes(self) -> Vec<u8> {
        self.data.to_vec()
    }

    /// Get the number of elements stored in the vector.
    #[inline]
    pub fn len(&self) -> usize {
        self.length
    }

    /// Returns `true` if the nibble vector has a length of 0.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Fetch a single entry from the vector.
    ///
    /// Guaranteed to be a value in the interval [0, 15].
    ///
    /// **Panics** if `idx >= self.len()`.
    #[inline]
    pub fn get(&self, idx: usize) -> u8 {
        if idx >= self.length {
            panic!(
                "attempted access beyond vector end. len is {}, index is {}",
                self.length, idx
            );
        }
        let vec_idx = idx / 2;
        match idx % 2 {
            // If the index is even, take the first (most significant) half of the stored byte.
            0 => self.data[vec_idx] >> 4,
            // If the index is odd, take the second (least significant) half.
            _ => self.data[vec_idx] & 0x0F,
        }
    }

    /// Add a single nibble to the vector.
    ///
    /// Only the 4 least-significant bits of the value are used.
    #[inline]
    pub fn push(&mut self, val: u8) {
        if self.length % 2 == 0 {
            self.data.push(val << 4);
        } else {
            let vec_len = self.data.len();

            // Zero the second half of the last byte just to be safe.
            self.data[vec_len - 1] &= 0xF0;

            // Write the new value.
            self.data[vec_len - 1] |= val & 0x0F;
        }
        self.length += 1;
    }

    /// Split the vector into two parts.
    ///
    /// All elements at or following the given index are returned in a new `NibbleVec`,
    /// with exactly `idx` elements remaining in this vector.
    ///
    /// **Panics** if `idx > self.len()`.
    pub fn split(&mut self, idx: usize) -> NibbleVec {
        if idx > self.length {
            panic!(
                "attempted to split past vector end. len is {}, index is {}",
                self.length, idx
            );
        } else if idx == self.length {
            NibbleVec::new()
        } else if idx % 2 == 0 {
            self.split_even(idx)
        } else {
            self.split_odd(idx)
        }
    }

    /// Split function for odd *indices*.
    #[inline]
    fn split_odd(&mut self, idx: usize) -> NibbleVec {
        let tail_vec_size = (self.length - idx) / 2;
        let mut tail = NibbleVec::from_small_vec(0);

        // Perform an overlap copy, copying the last nibble of the original vector only if
        // the length of the new tail is *odd*.
        let tail_length = self.length - idx;
        let take_last = tail_length % 2 == 1;
        self.overlap_copy(
            idx / 2,
            self.data.len(),
            &mut tail.data,
            &mut tail.length,
            take_last,
        );

        // Remove the copied bytes, being careful to skip the idx byte.
        for _ in (idx / 2 + 1)..self.data.len() {
            self.data.pop();
        }

        // Zero the second half of the index byte so as to maintain the last-nibble invariant.
        self.data[idx / 2] &= 0xF0;

        // Update the length of the first NibbleVec.
        self.length = idx;

        tail
    }

    /// Split function for even *indices*.
    #[inline]
    fn split_even(&mut self, idx: usize) -> NibbleVec {
        // Avoid allocating a temporary vector by copying all the bytes in order, then popping them.

        // Possible to prove: l_d - ⌊i / 2⌋ = ⌊(l_v - i + 1) / 2⌋
        //  where l_d = self.data.len()
        //        l_v = self.length
        let tail_vec_size = (self.length - idx + 1) / 2;
        let half_idx = idx / 2;
        let mut tail = NibbleVec::from_small_vec(0);

        // Copy the bytes.
        for i in half_idx..self.data.len() {
            tail.data.push(self.data[i]);
        }

        // Pop the same bytes.
        for _ in half_idx..self.data.len() {
            self.data.pop();
        }

        // Update lengths.
        tail.length = self.length - idx;
        self.length = idx;

        tail
    }

    /// Copy data between the second half of self.data[start] and
    /// self.data[end - 1]. The second half of the last entry is included
    /// if include_last is true.
    #[inline]
    fn overlap_copy(
        &self,
        start: usize,
        end: usize,
        vec: &mut SmallVec<[u8; 20]>,
        length: &mut usize,
        include_last: bool,
    ) {
        // Copy up to the first half of the last byte.
        for i in start..(end - 1) {
            // The first half is the second half of the old entry.
            let first_half = self.data[i] & 0x0f;

            // The second half is the first half of the next entry.
            let second_half = self.data[i + 1] >> 4;

            vec.push((first_half << 4) | second_half);
            *length += 2;
        }

        if include_last {
            let last = self.data[end - 1] & 0x0f;
            vec.push(last << 4);
            *length += 1;
        }
    }

    /// Append another nibble vector whilst consuming this vector.
    #[inline]
    pub fn join(mut self, other: &NibbleVec) -> NibbleVec {
        // If the length is even, we can append directly.
        if self.length % 2 == 0 {
            self.length += other.length;
            self.data.extend_from_slice(&other.data);
            return self;
        }

        // If the other vector is empty, bail out.
        if other.is_empty() {
            return self;
        }

        // If the length is odd, we have to perform an overlap copy.
        // Copy the first half of the first element, to make the vector an even length.
        self.push(other.get(0));

        // Copy the rest of the vector using an overlap copy.
        let take_last = other.len() % 2 == 0;
        other.overlap_copy(
            0,
            other.data.len(),
            &mut self.data,
            &mut self.length,
            take_last,
        );

        self
    }
}

impl PartialEq<NibbleVec> for NibbleVec {
    #[inline]
    fn eq(&self, other: &NibbleVec) -> bool {
        self.length == other.length && self.data == other.data
    }
}

impl Eq for NibbleVec {}

/// Compare a `NibbleVec` and a slice of bytes *element-by-element*.
/// Bytes are **not** interpreted as two `NibbleVec` entries.
impl PartialEq<[u8]> for NibbleVec {
    #[inline]
    fn eq(&self, other: &[u8]) -> bool {
        if other.len() != self.len() {
            return false;
        }

        for (i, x) in other.iter().enumerate() {
            if self.get(i) != *x {
                return false;
            }
        }
        true
    }
}

impl Debug for NibbleVec {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "NibbleVec [")?;

        if !self.is_empty() {
            write!(fmt, "{}", self.get(0))?;
        }

        for i in 1..self.len() {
            write!(fmt, ", {}", self.get(i))?;
        }
        write!(fmt, "]")
    }
}

impl From<Vec<u8>> for NibbleVec {
    #[inline]
    fn from(v: Vec<u8>) -> NibbleVec {
        NibbleVec::from_byte_vec(v)
    }
}

impl<'a> From<&'a [u8]> for NibbleVec {
    #[inline]
    fn from(v: &[u8]) -> NibbleVec {
        NibbleVec::from_byte_vec(v.into())
    }
}

impl Into<Vec<u8>> for NibbleVec {
    #[inline]
    fn into(self) -> Vec<u8> {
        self.data.to_vec()
    }
}

impl<'a> Into<Vec<u8>> for &'a NibbleVec {
    #[inline]
    fn into(self) -> Vec<u8> {
        self.data.to_vec()
    }
}
<|endoftext|>