<commit_before>// Third attempt.

use mio::{self, unix, Ready};
use mio::unix::UnixReady;
use std::io;
use dbus::{Connection, ConnectionItems, ConnectionItem, Watch, WatchEvent, Message, Error as DBusError};
use futures::{Async, Future, Stream, Poll};
use futures::sync::{oneshot, mpsc};
use tokio_core::reactor::{PollEvented, Handle as CoreHandle};
use std::rc::Rc;
use std::os::raw::c_uint;
use std::cell::RefCell;
use std::collections::HashMap;
use std::os::unix::io::RawFd;

type MCallMap = Rc<RefCell<HashMap<u32, oneshot::Sender<Message>>>>;

type MStream = Rc<RefCell<Option<mpsc::UnboundedSender<Message>>>>;

#[derive(Debug)]
/// A Tokio enabled D-Bus connection.
///
/// While an AConnection exists, it will consume all incoming messages.
/// Creating more than one AConnection for the same Connection is not recommended.
pub struct AConnection {
    conn: Rc<Connection>,
    quit: Option<oneshot::Sender<()>>,
    callmap: MCallMap,
    msgstream: MStream,
}

impl AConnection {
    /// Create an AConnection, which spawns a task on the core.
    ///
    /// The task handles incoming messages, and continues to do so until the
    /// AConnection is dropped.
    pub fn new(c: Rc<Connection>, h: CoreHandle) -> io::Result<AConnection> {
        let (tx, rx) = oneshot::channel();
        let map: MCallMap = Default::default();
        let istream: MStream = Default::default();
        let mut d = ADriver {
            conn: c.clone(),
            fds: HashMap::new(),
            core: h.clone(),
            quit: rx,
            callmap: map.clone(),
            msgstream: istream.clone(),
        };
        let i = AConnection {
            conn: c,
            quit: Some(tx),
            callmap: map,
            msgstream: istream,
        };
        for w in i.conn.watch_fds() { d.modify_watch(w)?; }
        h.spawn(d);
        Ok(i)
    }

    /// Sends a method call message, and returns a Future for the method return.
    pub fn method_call(&self, m: Message) -> Result<AMethodCall, &'static str> {
        let r = self.conn.send(m).map_err(|_| "D-Bus send error")?;
        let (tx, rx) = oneshot::channel();
        let mut map = self.callmap.borrow_mut();
        map.insert(r, tx); // TODO: error check for duplicate entries. Should not happen, but if it does...
        let mc = AMethodCall { serial: r, callmap: self.callmap.clone(), inner: rx };
        Ok(mc)
    }

    /// Returns a stream of incoming messages.
    ///
    /// Creating more than one stream for the same AConnection is not supported; this function will
    /// fail with an error if you try. Drop the first stream if you need to create a second one.
    pub fn messages(&self) -> Result<AMessageStream, &'static str> {
        let mut i = self.msgstream.borrow_mut();
        if i.is_some() { return Err("Another instance of AMessageStream already exists"); }
        let (tx, rx) = mpsc::unbounded();
        *i = Some(tx);
        Ok(AMessageStream { inner: rx, stream: self.msgstream.clone() })
    }
}

impl Drop for AConnection {
    fn drop(&mut self) { let _ = self.quit.take().unwrap().send(()); }
}

#[derive(Debug)]
// Internal struct; this is the future spawned on the core.
struct ADriver {
    conn: Rc<Connection>,
    fds: HashMap<RawFd, PollEvented<AWatch>>,
    core: CoreHandle,
    quit: oneshot::Receiver<()>,
    callmap: MCallMap,
    msgstream: MStream,
}

impl ADriver {
    fn modify_watch(&mut self, w: Watch) -> io::Result<()> {
        if !w.readable() && !w.writable() {
            self.fds.remove(&w.fd());
        } else {

            if let Some(evented) = self.fds.get(&w.fd()) {
                let ww = evented.get_ref().0;
                if ww.readable() == w.readable() || ww.readable() == w.readable() { return Ok(()) };
            }
            self.fds.remove(&w.fd());

            let z = PollEvented::new(AWatch(w), &self.core)?;

            self.fds.insert(w.fd(), z);
        }
        Ok(())
    }

    fn send_stream(&self, m: Message) {
        self.msgstream.borrow().as_ref().map(|z| { z.send(m).unwrap() });
    }

    fn handle_items(&mut self, items: ConnectionItems) {
        // TODO: What about all unwrapping in this function
        for i in items {
            // println!("Got: {:?}", i);
            match i {
                ConnectionItem::WatchFd(w) => { self.modify_watch(w).unwrap(); },
                ConnectionItem::MethodReturn(m) => {
                    let mut map = self.callmap.borrow_mut();
                    let serial = m.get_reply_serial().unwrap();
                    let r = map.remove(&serial);
                    if let Some(r) = r { r.send(m).unwrap(); }
                    else { self.send_stream(m) }
                }
                ConnectionItem::Nothing => return,
                ConnectionItem::Signal(m) => self.send_stream(m),
                ConnectionItem::MethodCall(m) => self.send_stream(m),
            }
        }
    }
}

impl Future for ADriver {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Result<Async<()>, ()> {
        let q = self.quit.poll();
        if q != Ok(Async::NotReady) { return Ok(Async::Ready(())); }

        let cc = self.conn.clone(); // Borrow checker made me do this
        let mut items = None;
        for w in self.fds.values() {
            let mut mask = UnixReady::hup() | UnixReady::error();
            if w.get_ref().0.readable() { mask = mask | Ready::readable().into(); }
            if w.get_ref().0.writable() { mask = mask | Ready::writable().into(); }
            let pr = w.poll_ready(*mask);
            // println!("{:?} is {:?}", w.get_ref().0.fd(), pr);
            let ur = if let Async::Ready(t) = pr { UnixReady::from(t) } else { continue };
            let flags =
                if ur.is_readable() { WatchEvent::Readable as c_uint } else { 0 } +
                if ur.is_writable() { WatchEvent::Writable as c_uint } else { 0 } +
                if ur.is_hup() { WatchEvent::Hangup as c_uint } else { 0 } +
                if ur.is_error() { WatchEvent::Error as c_uint } else { 0 };
            // println!("{:?} is {:?}", w.get_ref().0.fd(), ur);
            items = Some(cc.watch_handle(w.get_ref().0.fd(), flags));
            if ur.is_readable() { w.need_read() };
            if ur.is_writable() { w.need_write() };
        };
        self.handle_items(items.unwrap_or(cc.iter(0)));
        Ok(Async::NotReady)
    }
}

#[derive(Debug)]
struct AWatch(Watch);

impl mio::Evented for AWatch {
    fn register(&self,
                poll: &mio::Poll,
                token: mio::Token,
                mut interest: mio::Ready,
                opts: mio::PollOpt) -> io::Result<()>
    {
        if !self.0.readable() { interest.remove(mio::Ready::readable()) };
        if !self.0.writable() { interest.remove(mio::Ready::writable()) };
        unix::EventedFd(&self.0.fd()).register(poll, token, interest, opts)
    }

    fn reregister(&self,
                  poll: &mio::Poll,
                  token: mio::Token,
                  mut interest: mio::Ready,
                  opts: mio::PollOpt) -> io::Result<()>
    {
        if !self.0.readable() { interest.remove(mio::Ready::readable()) };
        if !self.0.writable() { interest.remove(mio::Ready::writable()) };
        unix::EventedFd(&self.0.fd()).reregister(poll, token, interest, opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        unix::EventedFd(&self.0.fd()).deregister(poll)
    }
}

#[derive(Debug)]
/// A Future that resolves when a method call is replied to.
pub struct AMethodCall {
    serial: u32,
    callmap: MCallMap,
    inner: oneshot::Receiver<Message>,
}

impl Future for AMethodCall {
    type Item = Message;
    type Error = DBusError;

    fn poll(&mut self) -> Result<Async<Self::Item>, Self::Error> {
        let x = self.inner.poll().map_err(|_| DBusError::new_custom("org.freedesktop.DBus.Failed", "Tokio cancelled future"))?;
        if let Async::Ready(mut m) = x {
            m.as_result()?;
            Ok(Async::Ready(m))
        } else { Ok(Async::NotReady) }
    }
}

impl Drop for AMethodCall {
    fn drop(&mut self) {
        let mut map = self.callmap.borrow_mut();
        map.remove(&self.serial);
    }
}

#[derive(Debug)]
/// A Stream of incoming messages.
///
/// Messages already processed (method returns for AMethodCall, fd changes)
/// are already consumed and will not be present in the stream.
pub struct AMessageStream {
    inner: mpsc::UnboundedReceiver<Message>,
    stream: MStream,
}

impl Stream for AMessageStream {
    type Item = Message;
    type Error = ();
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // println!("Polling message stream");
        let r = self.inner.poll();
        // println!("msgstream {:?}", r);
        r
    }
}

impl Drop for AMessageStream {
    fn drop(&mut self) {
        *self.stream.borrow_mut() = None;
    }
}

#[test]
fn aconnection_test() {
    let conn = Rc::new(Connection::get_private(::dbus::BusType::Session).unwrap());
    let mut core = ::tokio_core::reactor::Core::new().unwrap();
    let aconn = AConnection::new(conn.clone(), core.handle()).unwrap();

    let m = ::dbus::Message::new_method_call("org.freedesktop.DBus", "/", "org.freedesktop.DBus", "ListNames").unwrap();
    let reply = core.run(aconn.method_call(m).unwrap()).unwrap();
    let z: Vec<&str> = reply.get1().unwrap();
    println!("got reply: {:?}", z);
    assert!(z.iter().any(|v| *v == "org.freedesktop.DBus"));
}

#[test]
fn astream_test() {
    let conn = Rc::new(Connection::get_private(::dbus::BusType::Session).unwrap());
    let mut core = ::tokio_core::reactor::Core::new().unwrap();
    let aconn = AConnection::new(conn.clone(), core.handle()).unwrap();

    let items: AMessageStream = aconn.messages().unwrap();
    let signals = items.filter_map(|m| if m.msg_type() == ::dbus::MessageType::Signal { Some(m) } else { None });
    let firstsig = core.run(signals.into_future()).map(|(x, _)| x).map_err(|(x, _)| x).unwrap();
    println!("first signal was: {:?}", firstsig);
    assert_eq!(firstsig.unwrap().msg_type(), ::dbus::MessageType::Signal);
}

<commit_msg>Force libdbus polling to be level-triggered<commit_after>// Third attempt.

use mio::{self, unix, Ready};
use mio::unix::UnixReady;
use std::io;
use dbus::{Connection, ConnectionItems, ConnectionItem, Watch, WatchEvent, Message, Error as DBusError};
use futures::{Async, Future, Stream, Poll};
use futures::sync::{oneshot, mpsc};
use tokio_core::reactor::{PollEvented, Handle as CoreHandle};
use std::rc::Rc;
use std::os::raw::c_uint;
use std::cell::RefCell;
use std::collections::HashMap;
use std::os::unix::io::RawFd;

type MCallMap = Rc<RefCell<HashMap<u32, oneshot::Sender<Message>>>>;

type MStream = Rc<RefCell<Option<mpsc::UnboundedSender<Message>>>>;

#[derive(Debug)]
/// A Tokio enabled D-Bus connection.
///
/// While an AConnection exists, it will consume all incoming messages.
/// Creating more than one AConnection for the same Connection is not recommended.
pub struct AConnection {
    conn: Rc<Connection>,
    quit: Option<oneshot::Sender<()>>,
    callmap: MCallMap,
    msgstream: MStream,
}

impl AConnection {
    /// Create an AConnection, which spawns a task on the core.
    ///
    /// The task handles incoming messages, and continues to do so until the
    /// AConnection is dropped.
    pub fn new(c: Rc<Connection>, h: CoreHandle) -> io::Result<AConnection> {
        let (tx, rx) = oneshot::channel();
        let map: MCallMap = Default::default();
        let istream: MStream = Default::default();
        let mut d = ADriver {
            conn: c.clone(),
            fds: HashMap::new(),
            core: h.clone(),
            quit: rx,
            callmap: map.clone(),
            msgstream: istream.clone(),
        };
        let i = AConnection {
            conn: c,
            quit: Some(tx),
            callmap: map,
            msgstream: istream,
        };
        for w in i.conn.watch_fds() { d.modify_watch(w)?; }
        h.spawn(d);
        Ok(i)
    }

    /// Sends a method call message, and returns a Future for the method return.
    pub fn method_call(&self, m: Message) -> Result<AMethodCall, &'static str> {
        let r = self.conn.send(m).map_err(|_| "D-Bus send error")?;
        let (tx, rx) = oneshot::channel();
        let mut map = self.callmap.borrow_mut();
        map.insert(r, tx); // TODO: error check for duplicate entries. Should not happen, but if it does...
        let mc = AMethodCall { serial: r, callmap: self.callmap.clone(), inner: rx };
        Ok(mc)
    }

    /// Returns a stream of incoming messages.
    ///
    /// Creating more than one stream for the same AConnection is not supported; this function will
    /// fail with an error if you try. Drop the first stream if you need to create a second one.
    pub fn messages(&self) -> Result<AMessageStream, &'static str> {
        let mut i = self.msgstream.borrow_mut();
        if i.is_some() { return Err("Another instance of AMessageStream already exists"); }
        let (tx, rx) = mpsc::unbounded();
        *i = Some(tx);
        Ok(AMessageStream { inner: rx, stream: self.msgstream.clone() })
    }
}

impl Drop for AConnection {
    fn drop(&mut self) { let _ = self.quit.take().unwrap().send(()); }
}

#[derive(Debug)]
// Internal struct; this is the future spawned on the core.
struct ADriver {
    conn: Rc<Connection>,
    fds: HashMap<RawFd, PollEvented<AWatch>>,
    core: CoreHandle,
    quit: oneshot::Receiver<()>,
    callmap: MCallMap,
    msgstream: MStream,
}

impl ADriver {
    fn modify_watch(&mut self, w: Watch) -> io::Result<()> {
        if !w.readable() && !w.writable() {
            self.fds.remove(&w.fd());
        } else {

            if let Some(evented) = self.fds.get(&w.fd()) {
                let ww = evented.get_ref().0;
                if ww.readable() == w.readable() || ww.readable() == w.readable() { return Ok(()) };
            }
            self.fds.remove(&w.fd());

            let z = PollEvented::new(AWatch(w), &self.core)?;

            self.fds.insert(w.fd(), z);
        }
        Ok(())
    }

    fn send_stream(&self, m: Message) {
        self.msgstream.borrow().as_ref().map(|z| { z.send(m).unwrap() });
    }

    fn handle_items(&mut self, items: ConnectionItems) {
        // TODO: What about all unwrapping in this function
        for i in items {
            // println!("Got: {:?}", i);
            match i {
                ConnectionItem::WatchFd(w) => { self.modify_watch(w).unwrap(); },
                ConnectionItem::MethodReturn(m) => {
                    let mut map = self.callmap.borrow_mut();
                    let serial = m.get_reply_serial().unwrap();
                    let r = map.remove(&serial);
                    if let Some(r) = r { r.send(m).unwrap(); }
                    else { self.send_stream(m) }
                }
                ConnectionItem::Nothing => return,
                ConnectionItem::Signal(m) => self.send_stream(m),
                ConnectionItem::MethodCall(m) => self.send_stream(m),
            }
        }
    }
}

impl Future for ADriver {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Result<Async<()>, ()> {
        let q = self.quit.poll();
        if q != Ok(Async::NotReady) { return Ok(Async::Ready(())); }

        let cc = self.conn.clone(); // Borrow checker made me do this
        let mut items = None;
        for w in self.fds.values() {
            let mut mask = UnixReady::hup() | UnixReady::error();
            if w.get_ref().0.readable() { mask = mask | Ready::readable().into(); }
            if w.get_ref().0.writable() { mask = mask | Ready::writable().into(); }
            let pr = w.poll_ready(*mask);
            // println!("{:?} is {:?}", w.get_ref().0.fd(), pr);
            let ur = if let Async::Ready(t) = pr { UnixReady::from(t) } else { continue };
            let flags =
                if ur.is_readable() { WatchEvent::Readable as c_uint } else { 0 } +
                if ur.is_writable() { WatchEvent::Writable as c_uint } else { 0 } +
                if ur.is_hup() { WatchEvent::Hangup as c_uint } else { 0 } +
                if ur.is_error() { WatchEvent::Error as c_uint } else { 0 };
            // println!("{:?} is {:?}", w.get_ref().0.fd(), ur);
            items = Some(cc.watch_handle(w.get_ref().0.fd(), flags));
            if ur.is_readable() { w.need_read() };
            if ur.is_writable() { w.need_write() };
        };
        self.handle_items(items.unwrap_or(cc.iter(0)));
        Ok(Async::NotReady)
    }
}

#[derive(Debug)]
struct AWatch(Watch);

impl mio::Evented for AWatch {
    fn register(&self,
                poll: &mio::Poll,
                token: mio::Token,
                mut interest: mio::Ready,
                mut opts: mio::PollOpt) -> io::Result<()>
    {
        if !self.0.readable() { interest.remove(mio::Ready::readable()) };
        if !self.0.writable() { interest.remove(mio::Ready::writable()) };
        opts.remove(mio::PollOpt::edge());
        opts.insert(mio::PollOpt::level());
        unix::EventedFd(&self.0.fd()).register(poll, token, interest, opts)
    }

    fn reregister(&self,
                  poll: &mio::Poll,
                  token: mio::Token,
                  mut interest: mio::Ready,
                  mut opts: mio::PollOpt) -> io::Result<()>
    {
        if !self.0.readable() { interest.remove(mio::Ready::readable()) };
        if !self.0.writable() { interest.remove(mio::Ready::writable()) };
        opts.remove(mio::PollOpt::edge());
        opts.insert(mio::PollOpt::level());
        unix::EventedFd(&self.0.fd()).reregister(poll, token, interest, opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        unix::EventedFd(&self.0.fd()).deregister(poll)
    }
}

#[derive(Debug)]
/// A Future that resolves when a method call is replied to.
pub struct AMethodCall {
    serial: u32,
    callmap: MCallMap,
    inner: oneshot::Receiver<Message>,
}

impl Future for AMethodCall {
    type Item = Message;
    type Error = DBusError;

    fn poll(&mut self) -> Result<Async<Self::Item>, Self::Error> {
        let x = self.inner.poll().map_err(|_| DBusError::new_custom("org.freedesktop.DBus.Failed", "Tokio cancelled future"))?;
        if let Async::Ready(mut m) = x {
            m.as_result()?;
            Ok(Async::Ready(m))
        } else { Ok(Async::NotReady) }
    }
}

impl Drop for AMethodCall {
    fn drop(&mut self) {
        let mut map = self.callmap.borrow_mut();
        map.remove(&self.serial);
    }
}

#[derive(Debug)]
/// A Stream of incoming messages.
///
/// Messages already processed (method returns for AMethodCall, fd changes)
/// are already consumed and will not be present in the stream.
pub struct AMessageStream {
    inner: mpsc::UnboundedReceiver<Message>,
    stream: MStream,
}

impl Stream for AMessageStream {
    type Item = Message;
    type Error = ();
    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        // println!("Polling message stream");
        let r = self.inner.poll();
        // println!("msgstream {:?}", r);
        r
    }
}

impl Drop for AMessageStream {
    fn drop(&mut self) {
        *self.stream.borrow_mut() = None;
    }
}

#[test]
fn aconnection_test() {
    let conn = Rc::new(Connection::get_private(::dbus::BusType::Session).unwrap());
    let mut core = ::tokio_core::reactor::Core::new().unwrap();
    let aconn = AConnection::new(conn.clone(), core.handle()).unwrap();

    let m = ::dbus::Message::new_method_call("org.freedesktop.DBus", "/", "org.freedesktop.DBus", "ListNames").unwrap();
    let reply = core.run(aconn.method_call(m).unwrap()).unwrap();
    let z: Vec<&str> = reply.get1().unwrap();
    println!("got reply: {:?}", z);
    assert!(z.iter().any(|v| *v == "org.freedesktop.DBus"));
}

#[test]
fn astream_test() {
    let conn = Rc::new(Connection::get_private(::dbus::BusType::Session).unwrap());
    let mut core = ::tokio_core::reactor::Core::new().unwrap();
    let aconn = AConnection::new(conn.clone(), core.handle()).unwrap();

    let items: AMessageStream = aconn.messages().unwrap();
    let signals = items.filter_map(|m| if m.msg_type() == ::dbus::MessageType::Signal { Some(m) } else { None });
    let firstsig = core.run(signals.into_future()).map(|(x, _)| x).map_err(|(x, _)| x).unwrap();
    println!("first signal was: {:?}", firstsig);
    assert_eq!(firstsig.unwrap().msg_type(), ::dbus::MessageType::Signal);
}

<|endoftext|>


<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate libc;

use std::fmt;

pub mod validators;

/// Constant to convert seconds to nanoseconds.
pub const NANOS_PER_SECOND: u64 = 1_000_000_000;

/// Wrapper over `libc::clockid_t` to specify Linux Kernel clock source.
pub enum ClockType {
    /// Equivalent to `libc::CLOCK_MONOTONIC`.
    Monotonic,
    /// Equivalent to `libc::CLOCK_REALTIME`.
    Real,
    /// Equivalent to `libc::CLOCK_PROCESS_CPUTIME_ID`.
    ProcessCpu,
    /// Equivalent to `libc::CLOCK_THREAD_CPUTIME_ID`.
    ThreadCpu,
}

impl Into<libc::clockid_t> for ClockType {
    fn into(self) -> libc::clockid_t {
        match self {
            ClockType::Monotonic => libc::CLOCK_MONOTONIC,
            ClockType::Real => libc::CLOCK_REALTIME,
            ClockType::ProcessCpu => libc::CLOCK_PROCESS_CPUTIME_ID,
            ClockType::ThreadCpu => libc::CLOCK_THREAD_CPUTIME_ID,
        }
    }
}

/// Structure representing the date in local time with nanosecond precision.
pub struct LocalTime {
    /// Seconds in current minute.
    sec: i32,
    /// Minutes in current hour.
    min: i32,
    /// Hours in current day, 24H format.
    hour: i32,
    /// Days in current month.
    mday: i32,
    /// Months in current year.
    mon: i32,
    /// Years passed since 1900 BC.
    year: i32,
    /// Nanoseconds in current second.
    nsec: i64,
}

impl LocalTime {
    /// Returns the [LocalTime](struct.LocalTime.html) structure for the calling moment.
    pub fn now() -> LocalTime {
        let mut timespec = libc::timespec {
            tv_sec: 0,
            tv_nsec: 0,
        };
        let mut tm: libc::tm = libc::tm {
            tm_sec: 0,
            tm_min: 0,
            tm_hour: 0,
            tm_mday: 0,
            tm_mon: 0,
            tm_year: 0,
            tm_wday: 0,
            tm_yday: 0,
            tm_isdst: 0,
            tm_gmtoff: 0,
            tm_zone: std::ptr::null(),
        };

        // Safe because the parameters are valid.
        unsafe {
            libc::clock_gettime(libc::CLOCK_REALTIME, &mut timespec);
            libc::localtime_r(&timespec.tv_sec, &mut tm);
        }

        LocalTime {
            sec: tm.tm_sec,
            min: tm.tm_min,
            hour: tm.tm_hour,
            mday: tm.tm_mday,
            mon: tm.tm_mon,
            year: tm.tm_year,
            nsec: timespec.tv_nsec,
        }
    }
}

impl fmt::Display for LocalTime {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}-{:02}-{:02}T{:02}:{:02}:{:02}.{:09}",
            self.year + 1900,
            self.mon,
            self.mday,
            self.hour,
            self.min,
            self.sec,
            self.nsec
        )
    }
}

/// Returns a timestamp in nanoseconds from a monotonic clock.
///
/// Uses `_rdstc` on `x86_64` and [`get_time`](fn.get_time.html) on other architectures.
pub fn timestamp_cycles() -> u64 {
    #[cfg(target_arch = "x86_64")]
    // Safe because there's nothing that can go wrong with this call.
    unsafe {
        std::arch::x86_64::_rdtsc() as u64
    }
    #[cfg(not(target_arch = "x86_64"))]
    {
        get_time(ClockType::Monotonic)
    }
}

/// Returns a timestamp in nanoseconds based on the provided clock type.
///
/// # Arguments
///
/// * `clock_type` - Identifier of the Linux Kernel clock on which to act.
pub fn get_time(clock_type: ClockType) -> u64 {
    let mut time_struct = libc::timespec {
        tv_sec: 0,
        tv_nsec: 0,
    };
    // Safe because the parameters are valid.
    unsafe { libc::clock_gettime(clock_type.into(), &mut time_struct) };
    seconds_to_nanoseconds(time_struct.tv_sec).unwrap() as u64 + (time_struct.tv_nsec as u64)
}

/// Converts a timestamp in seconds to an equivalent one in nanoseconds.
/// Returns `None` if the conversion overflows.
///
/// # Arguments
///
/// * `value` - Timestamp in seconds.
pub fn seconds_to_nanoseconds(value: i64) -> Option<i64> {
    value.checked_mul(NANOS_PER_SECOND as i64)
}

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

    #[test]
    fn test_get_time() {
        for _ in 0..1000 {
            assert!(get_time(ClockType::Monotonic) <= get_time(ClockType::Monotonic));
        }

        for _ in 0..1000 {
            assert!(get_time(ClockType::ProcessCpu) <= get_time(ClockType::ProcessCpu));
        }

        for _ in 0..1000 {
            assert!(get_time(ClockType::ThreadCpu) <= get_time(ClockType::ThreadCpu));
        }

        assert_ne!(get_time(ClockType::Real), 0);
    }

    #[test]
    fn test_local_time_display() {
        let local_time = LocalTime {
            sec: 30,
            min: 15,
            hour: 10,
            mday: 4,
            mon: 7,
            year: 119,
            nsec: 123_456_789,
        };
        assert_eq!(
            String::from("2019-07-04T10:15:30.123456789").cmp(&local_time.to_string()),
            Ordering::Equal
        );

        let local_time = LocalTime {
            sec: 5,
            min: 5,
            hour: 5,
            mday: 23,
            mon: 8,
            year: 44,
            nsec: 123,
        };
        assert_eq!(
            String::from("1944-08-23T05:05:05.000000123"),
            local_time.to_string()
        );
    }

    #[test]
    fn test_seconds_to_nanos() {
        assert_eq!(
            seconds_to_nanoseconds(100).unwrap() as u64,
            100 * NANOS_PER_SECOND
        );

        assert!(seconds_to_nanoseconds(9_223_372_037).is_none());
    }
}
<commit_msg>Fixed logged month being one less than actual<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

extern crate libc;

use std::fmt;

pub mod validators;

/// Constant to convert seconds to nanoseconds.
pub const NANOS_PER_SECOND: u64 = 1_000_000_000;

/// Wrapper over `libc::clockid_t` to specify Linux Kernel clock source.
pub enum ClockType {
    /// Equivalent to `libc::CLOCK_MONOTONIC`.
    Monotonic,
    /// Equivalent to `libc::CLOCK_REALTIME`.
    Real,
    /// Equivalent to `libc::CLOCK_PROCESS_CPUTIME_ID`.
    ProcessCpu,
    /// Equivalent to `libc::CLOCK_THREAD_CPUTIME_ID`.
    ThreadCpu,
}

impl Into<libc::clockid_t> for ClockType {
    fn into(self) -> libc::clockid_t {
        match self {
            ClockType::Monotonic => libc::CLOCK_MONOTONIC,
            ClockType::Real => libc::CLOCK_REALTIME,
            ClockType::ProcessCpu => libc::CLOCK_PROCESS_CPUTIME_ID,
            ClockType::ThreadCpu => libc::CLOCK_THREAD_CPUTIME_ID,
        }
    }
}

/// Structure representing the date in local time with nanosecond precision.
pub struct LocalTime {
    /// Seconds in current minute.
    sec: i32,
    /// Minutes in current hour.
    min: i32,
    /// Hours in current day, 24H format.
    hour: i32,
    /// Days in current month.
    mday: i32,
    /// Months in current year.
    mon: i32,
    /// Years passed since 1900 BC.
    year: i32,
    /// Nanoseconds in current second.
    nsec: i64,
}

impl LocalTime {
    /// Returns the [LocalTime](struct.LocalTime.html) structure for the calling moment.
    pub fn now() -> LocalTime {
        let mut timespec = libc::timespec {
            tv_sec: 0,
            tv_nsec: 0,
        };
        let mut tm: libc::tm = libc::tm {
            tm_sec: 0,
            tm_min: 0,
            tm_hour: 0,
            tm_mday: 0,
            tm_mon: 0,
            tm_year: 0,
            tm_wday: 0,
            tm_yday: 0,
            tm_isdst: 0,
            tm_gmtoff: 0,
            tm_zone: std::ptr::null(),
        };

        // Safe because the parameters are valid.
        unsafe {
            libc::clock_gettime(libc::CLOCK_REALTIME, &mut timespec);
            libc::localtime_r(&timespec.tv_sec, &mut tm);
        }

        LocalTime {
            sec: tm.tm_sec,
            min: tm.tm_min,
            hour: tm.tm_hour,
            mday: tm.tm_mday,
            mon: tm.tm_mon,
            year: tm.tm_year,
            nsec: timespec.tv_nsec,
        }
    }
}

impl fmt::Display for LocalTime {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}-{:02}-{:02}T{:02}:{:02}:{:02}.{:09}",
            self.year + 1900,
            self.mon + 1,
            self.mday,
            self.hour,
            self.min,
            self.sec,
            self.nsec
        )
    }
}

/// Returns a timestamp in nanoseconds from a monotonic clock.
///
/// Uses `_rdstc` on `x86_64` and [`get_time`](fn.get_time.html) on other architectures.
pub fn timestamp_cycles() -> u64 {
    #[cfg(target_arch = "x86_64")]
    // Safe because there's nothing that can go wrong with this call.
    unsafe {
        std::arch::x86_64::_rdtsc() as u64
    }
    #[cfg(not(target_arch = "x86_64"))]
    {
        get_time(ClockType::Monotonic)
    }
}

/// Returns a timestamp in nanoseconds based on the provided clock type.
///
/// # Arguments
///
/// * `clock_type` - Identifier of the Linux Kernel clock on which to act.
pub fn get_time(clock_type: ClockType) -> u64 {
    let mut time_struct = libc::timespec {
        tv_sec: 0,
        tv_nsec: 0,
    };
    // Safe because the parameters are valid.
    unsafe { libc::clock_gettime(clock_type.into(), &mut time_struct) };
    seconds_to_nanoseconds(time_struct.tv_sec).unwrap() as u64 + (time_struct.tv_nsec as u64)
}

/// Converts a timestamp in seconds to an equivalent one in nanoseconds.
/// Returns `None` if the conversion overflows.
///
/// # Arguments
///
/// * `value` - Timestamp in seconds.
pub fn seconds_to_nanoseconds(value: i64) -> Option<i64> {
    value.checked_mul(NANOS_PER_SECOND as i64)
}

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

    #[test]
    fn test_get_time() {
        for _ in 0..1000 {
            assert!(get_time(ClockType::Monotonic) <= get_time(ClockType::Monotonic));
        }

        for _ in 0..1000 {
            assert!(get_time(ClockType::ProcessCpu) <= get_time(ClockType::ProcessCpu));
        }

        for _ in 0..1000 {
            assert!(get_time(ClockType::ThreadCpu) <= get_time(ClockType::ThreadCpu));
        }

        assert_ne!(get_time(ClockType::Real), 0);
    }

    #[test]
    fn test_local_time_display() {
        let local_time = LocalTime {
            sec: 30,
            min: 15,
            hour: 10,
            mday: 4,
            mon: 6,
            year: 119,
            nsec: 123_456_789,
        };
        assert_eq!(
            String::from("2019-07-04T10:15:30.123456789"),
            local_time.to_string()
        );

        let local_time = LocalTime {
            sec: 5,
            min: 5,
            hour: 5,
            mday: 23,
            mon: 7,
            year: 44,
            nsec: 123,
        };
        assert_eq!(
            String::from("1944-08-23T05:05:05.000000123"),
            local_time.to_string()
        );

        let local_time = LocalTime::now();
        assert!(local_time.mon >= 0 && local_time.mon <= 11);
    }

    #[test]
    fn test_seconds_to_nanoseconds() {
        assert_eq!(
            seconds_to_nanoseconds(100).unwrap() as u64,
            100 * NANOS_PER_SECOND
        );

        assert!(seconds_to_nanoseconds(9_223_372_037).is_none());
    }
}
<|endoftext|>


<commit_before>use super::*;

use std::cmp::{self, Ordering};

pub(crate) fn prefix_encode(prefix: &[u8], buf: &[u8]) -> IVec {
    assert!(
        prefix <= buf,
        "prefix {:?} must be lexicographically <= to the encoded buf {:?}",
        prefix,
        buf
    );

    let max = u8::max_value() as usize;
    let zip = prefix.iter().zip(buf);
    let prefix_len = zip.take(max).take_while(|(a, b)| a == b).count();

    let encoded_len = 1 + buf.len() - prefix_len;

    let mut ret = Vec::with_capacity(encoded_len);
    ret.push(prefix_len as u8);
    ret.extend_from_slice(&buf[prefix_len..]);

    IVec::from(ret)
}

pub(crate) fn prefix_decode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(!buf.is_empty());

    let prefix_len = buf[0] as usize;
    let mut ret = Vec::with_capacity(prefix_len + buf.len() - 1);

    ret.extend_from_slice(&prefix[..prefix_len]);
    ret.extend_from_slice(&buf[1..]);

    ret
}

pub(crate) fn prefix_reencode(
    old_prefix: &[u8],
    new_prefix: &[u8],
    buf: &[u8],
) -> IVec {
    assert!(!buf.is_empty());

    let prefix_len = buf[0] as usize;
    let old_prefix = &old_prefix[..prefix_len];
    let buf = &buf[1..];

    let mut output = Vec::with_capacity(old_prefix.len() + buf.len());

    output.extend_from_slice(old_prefix);
    output.extend_from_slice(buf);

    assert!(
        new_prefix <= &output,
        "new prefix {:?} must be lexicographically <= to the encoded buf {:?}",
        new_prefix,
        buf
    );

    let max = u8::max_value() as usize;
    let zip = new_prefix.iter().zip(&output);
    let prefix_len = zip.take(max).take_while(|(a, b)| a == b).count();

    let skip = prefix_len - cmp::min(old_prefix.len(), max);
    let encoded_len = 1 + output.len() - prefix_len;

    output.clear();
    output.reserve(encoded_len);

    output.push(prefix_len as u8);
    output.extend_from_slice(&buf[skip..]);

    IVec::from(output)
}

// NB: the correctness of this function depends on
// the invariant that the prefix is ALWAYS lexicographically
// Less than or Equal to the keys that have been encoded
// using it. Otherwise this comparison would make no sense.
pub(crate) fn prefix_cmp(a: &[u8], b: &[u8]) -> Ordering {
    if a.is_empty() && b.is_empty() {
        return Ordering::Equal;
    } else if a.is_empty() && !b.is_empty() {
        return Ordering::Less;
    } else if !a.is_empty() && b.is_empty() {
        return Ordering::Greater;
    }

    if a[0] > b[0] {
        Ordering::Less
    } else if a[0] < b[0] {
        Ordering::Greater
    } else {
        a[1..].cmp(&b[1..])
    }
}

/// Compare `a` and `b`, assuming that `a` is prefix encoded and `b` is not.
pub(crate) fn prefix_cmp_encoded(
    a: &[u8],
    mut b: &[u8],
    mut prefix: &[u8],
) -> Ordering {
    assert!(!a.is_empty() && a[0] as usize <= prefix.len());

    let mut a_prefix_len = a[0];
    let a_suffix = &a[1..];

    while a_prefix_len > 0 {
        if b.is_empty() || prefix[0] > b[0] {
            return Ordering::Greater;
        } else if prefix[0] < b[0] {
            return Ordering::Less;
        }

        a_prefix_len -= 1;
        b = &b[1..];
        prefix = &prefix[1..];
    }

    a_suffix.cmp(b)
}

#[test]
fn test_prefix() {
    let prefix = b"cat";
    assert_eq!(prefix_encode(prefix, prefix), vec![prefix.len() as u8]);
    assert_eq!(prefix_encode(prefix, b"catt"), vec![3, b't']);
    assert_eq!(prefix_encode(prefix, b"cb"), vec![1, b'b']);
    assert_eq!(prefix_encode(prefix, b"caz"), vec![2, b'z']);
    assert_eq!(prefix_encode(prefix, b"cvar"), vec![1, b'v', b'a', b'r']);
    assert_eq!(prefix_encode(prefix, b"zig"), vec![0, b'z', b'i', b'g']);

    let prefix = b"";
    assert_eq!(prefix_encode(prefix, prefix), vec![prefix.len() as u8]);
    assert_eq!(prefix_encode(prefix, b"ca"), vec![0, b'c', b'a']);
    assert_eq!(prefix_encode(prefix, b"cat"), vec![0, b'c', b'a', b't']);
    assert_eq!(prefix_encode(prefix, b"cab"), vec![0, b'c', b'a', b'b']);
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![0, b'c', b'v', b'a', b'r']
    );
    assert_eq!(prefix_encode(prefix, b"zig"), vec![0, b'z', b'i', b'g']);

    let rtt = vec![b"" as &[u8], b"\x00cat", b"\x00", b"oyunwytounw\x00"];
    for item in rtt {
        assert_eq!(
            prefix_decode(prefix, &*prefix_encode(prefix, item)),
            item.to_vec()
        );
    }
}

#[test]
fn test_prefix_cmp() {
    assert_eq!(prefix_cmp(&[], &[]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[0], &[]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[], &[0]), Ordering::Less);

    assert_eq!(prefix_cmp(&[3], &[4]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[4, 3], &[3, 4]), Ordering::Less);

    assert_eq!(prefix_cmp(&[1], &[1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 3], &[1, 1]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 3]), Ordering::Less);
}

#[test]
fn test_prefix_cmp_encoded() {
    fn assert_pce(a: &[u8], b: &[u8], prefix: &[u8], expected: Ordering) {
        assert_eq!(
            prefix_cmp_encoded(a, &prefix_decode(prefix, b), prefix),
            expected
        );
    }

    let prefix = &vec![3, 3, 3, 3];
    assert_pce(&[0], &[0], prefix, Ordering::Equal);
    assert_pce(&[1], &[0], prefix, Ordering::Greater);
    assert_pce(&[0], &[1], prefix, Ordering::Less);

    assert_pce(&[3], &[2], prefix, Ordering::Greater);
    assert_pce(&[4, 3], &[3, 4], prefix, Ordering::Less);

    assert_pce(&[1], &[1], prefix, Ordering::Equal);
    assert_pce(&[1, 3, 3, 1], &[3, 1], prefix, Ordering::Equal);
    assert_pce(&[1, 3], &[1, 1], prefix, Ordering::Greater);
    assert_pce(&[1, 1], &[3, 3], prefix, Ordering::Less);
}
<commit_msg>avoid clippy warning on reference usage<commit_after>use super::*;

use std::cmp::{self, Ordering};

pub(crate) fn prefix_encode(prefix: &[u8], buf: &[u8]) -> IVec {
    assert!(
        prefix <= buf,
        "prefix {:?} must be lexicographically <= to the encoded buf {:?}",
        prefix,
        buf
    );

    let max = u8::max_value() as usize;
    let zip = prefix.iter().zip(buf);
    let prefix_len = zip.take(max).take_while(|(a, b)| a == b).count();

    let encoded_len = 1 + buf.len() - prefix_len;

    let mut ret = Vec::with_capacity(encoded_len);
    ret.push(prefix_len as u8);
    ret.extend_from_slice(&buf[prefix_len..]);

    IVec::from(ret)
}

pub(crate) fn prefix_decode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(!buf.is_empty());

    let prefix_len = buf[0] as usize;
    let mut ret = Vec::with_capacity(prefix_len + buf.len() - 1);

    ret.extend_from_slice(&prefix[..prefix_len]);
    ret.extend_from_slice(&buf[1..]);

    ret
}

pub(crate) fn prefix_reencode(
    old_prefix: &[u8],
    new_prefix: &[u8],
    buf: &[u8],
) -> IVec {
    assert!(!buf.is_empty());

    let prefix_len = buf[0] as usize;
    let old_prefix = &old_prefix[..prefix_len];
    let buf = &buf[1..];

    let mut output = Vec::with_capacity(old_prefix.len() + buf.len());

    output.extend_from_slice(old_prefix);
    output.extend_from_slice(buf);

    assert!(
        *new_prefix <= *output,
        "new prefix {:?} must be lexicographically <= to the encoded buf {:?}",
        new_prefix,
        buf
    );

    let max = u8::max_value() as usize;
    let zip = new_prefix.iter().zip(&output);
    let prefix_len = zip.take(max).take_while(|(a, b)| a == b).count();

    let skip = prefix_len - cmp::min(old_prefix.len(), max);
    let encoded_len = 1 + output.len() - prefix_len;

    output.clear();
    output.reserve(encoded_len);

    output.push(prefix_len as u8);
    output.extend_from_slice(&buf[skip..]);

    IVec::from(output)
}

// NB: the correctness of this function depends on
// the invariant that the prefix is ALWAYS lexicographically
// Less than or Equal to the keys that have been encoded
// using it. Otherwise this comparison would make no sense.
pub(crate) fn prefix_cmp(a: &[u8], b: &[u8]) -> Ordering {
    if a.is_empty() && b.is_empty() {
        return Ordering::Equal;
    } else if a.is_empty() && !b.is_empty() {
        return Ordering::Less;
    } else if !a.is_empty() && b.is_empty() {
        return Ordering::Greater;
    }

    if a[0] > b[0] {
        Ordering::Less
    } else if a[0] < b[0] {
        Ordering::Greater
    } else {
        a[1..].cmp(&b[1..])
    }
}

/// Compare `a` and `b`, assuming that `a` is prefix encoded and `b` is not.
pub(crate) fn prefix_cmp_encoded(
    a: &[u8],
    mut b: &[u8],
    mut prefix: &[u8],
) -> Ordering {
    assert!(!a.is_empty() && a[0] as usize <= prefix.len());

    let mut a_prefix_len = a[0];
    let a_suffix = &a[1..];

    while a_prefix_len > 0 {
        if b.is_empty() || prefix[0] > b[0] {
            return Ordering::Greater;
        } else if prefix[0] < b[0] {
            return Ordering::Less;
        }

        a_prefix_len -= 1;
        b = &b[1..];
        prefix = &prefix[1..];
    }

    a_suffix.cmp(b)
}

#[test]
fn test_prefix() {
    let prefix = b"cat";
    assert_eq!(prefix_encode(prefix, prefix), vec![prefix.len() as u8]);
    assert_eq!(prefix_encode(prefix, b"catt"), vec![3, b't']);
    assert_eq!(prefix_encode(prefix, b"cb"), vec![1, b'b']);
    assert_eq!(prefix_encode(prefix, b"caz"), vec![2, b'z']);
    assert_eq!(prefix_encode(prefix, b"cvar"), vec![1, b'v', b'a', b'r']);
    assert_eq!(prefix_encode(prefix, b"zig"), vec![0, b'z', b'i', b'g']);

    let prefix = b"";
    assert_eq!(prefix_encode(prefix, prefix), vec![prefix.len() as u8]);
    assert_eq!(prefix_encode(prefix, b"ca"), vec![0, b'c', b'a']);
    assert_eq!(prefix_encode(prefix, b"cat"), vec![0, b'c', b'a', b't']);
    assert_eq!(prefix_encode(prefix, b"cab"), vec![0, b'c', b'a', b'b']);
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![0, b'c', b'v', b'a', b'r']
    );
    assert_eq!(prefix_encode(prefix, b"zig"), vec![0, b'z', b'i', b'g']);

    let rtt = vec![b"" as &[u8], b"\x00cat", b"\x00", b"oyunwytounw\x00"];
    for item in rtt {
        assert_eq!(
            prefix_decode(prefix, &*prefix_encode(prefix, item)),
            item.to_vec()
        );
    }
}

#[test]
fn test_prefix_cmp() {
    assert_eq!(prefix_cmp(&[], &[]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[0], &[]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[], &[0]), Ordering::Less);

    assert_eq!(prefix_cmp(&[3], &[4]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[4, 3], &[3, 4]), Ordering::Less);

    assert_eq!(prefix_cmp(&[1], &[1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 3], &[1, 1]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 3]), Ordering::Less);
}

#[test]
fn test_prefix_cmp_encoded() {
    fn assert_pce(a: &[u8], b: &[u8], prefix: &[u8], expected: Ordering) {
        assert_eq!(
            prefix_cmp_encoded(a, &prefix_decode(prefix, b), prefix),
            expected
        );
    }

    let prefix = &vec![3, 3, 3, 3];
    assert_pce(&[0], &[0], prefix, Ordering::Equal);
    assert_pce(&[1], &[0], prefix, Ordering::Greater);
    assert_pce(&[0], &[1], prefix, Ordering::Less);

    assert_pce(&[3], &[2], prefix, Ordering::Greater);
    assert_pce(&[4, 3], &[3, 4], prefix, Ordering::Less);

    assert_pce(&[1], &[1], prefix, Ordering::Equal);
    assert_pce(&[1, 3, 3, 1], &[3, 1], prefix, Ordering::Equal);
    assert_pce(&[1, 3], &[1, 1], prefix, Ordering::Greater);
    assert_pce(&[1, 1], &[3, 3], prefix, Ordering::Less);
}
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::cmp;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::mem;
use std::ptr;
use std::ops::{Deref, DerefMut};
use std::str;
use std::str::Utf8Error;
use std::slice;

use array::{Array, ArrayExt};
use array::Index;
use CapacityError;
use errors::PubCrateNew;
use odds::char::encode_utf8;

#[cfg(feature="serde-1")]
use serde::{Serialize, Deserialize, Serializer, Deserializer};

/// A string with a fixed capacity.
///
/// The `ArrayString` is a string backed by a fixed size array. It keeps track
/// of its length.
///
/// The string is a contiguous value that you can store directly on the stack
/// if needed.
#[derive(Copy)]
pub struct ArrayString<A: Array<Item=u8>> {
    xs: A,
    len: A::Index,
}

impl<A: Array<Item=u8>> ArrayString<A> {
    /// Create a new empty `ArrayString`.
    ///
    /// Capacity is inferred from the type parameter.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 16]>::new();
    /// string.push_str("foo");
    /// assert_eq!(&string[..], "foo");
    /// assert_eq!(string.capacity(), 16);
    /// ```
    pub fn new() -> ArrayString<A> {
        unsafe {
            ArrayString {
                xs: ::new_array(),
                len: Index::from(0),
            }
        }
    }

    /// Create a new `ArrayString` from a `str`.
    ///
    /// Capacity is inferred from the type parameter.
    ///
    /// **Errors** if the backing array is not large enough to fit the string.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 3]>::from("foo").unwrap();
    /// assert_eq!(&string[..], "foo");
    /// assert_eq!(string.len(), 3);
    /// assert_eq!(string.capacity(), 3);
    /// ```
    pub fn from(s: &str) -> Result<Self, CapacityError<&str>> {
        let mut arraystr = Self::new();
        arraystr.try_push_str(s)?;
        Ok(arraystr)
    }

    /// Create a new `ArrayString` from a byte string literal.
    ///
    /// **Errors** if the byte string literal is not valid UTF-8.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let string = ArrayString::from_byte_string(b"hello world").unwrap();
    /// ```
    pub fn from_byte_string(b: &A) -> Result<Self, Utf8Error> {
        let mut arraystr = Self::new();
        let s = try!(str::from_utf8(b.as_slice()));
        let _result = arraystr.try_push_str(s);
        debug_assert!(_result.is_ok());
        Ok(arraystr)
    }

    /// Return the capacity of the `ArrayString`.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let string = ArrayString::<[_; 3]>::new();
    /// assert_eq!(string.capacity(), 3);
    /// ```
    #[inline]
    pub fn capacity(&self) -> usize { A::capacity() }

    /// Return if the `ArrayString` is completely filled.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 1]>::new();
    /// assert!(!string.is_full());
    /// string.push_str("A");
    /// assert!(string.is_full());
    /// ```
    pub fn is_full(&self) -> bool { self.len() == self.capacity() }

    /// Adds the given char to the end of the string.
    ///
    /// Returns `Ok` if the push succeeds.
    ///
    /// ***Panics*** if the backing array is not large enough to fit the additional char.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 2]>::new();
    ///
    /// string.push('a');
    /// string.push('b');
    ///
    /// assert_eq!(&string[..], "ab");
    /// ```
    pub fn push(&mut self, c: char) {
        self.try_push(c).unwrap();
    }

    /// Adds the given char to the end of the string.
    ///
    /// Returns `Ok` if the push succeeds.
    ///
    /// **Errors** if the backing array is not large enough to fit the additional char.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 2]>::new();
    ///
    /// string.try_push('a').unwrap();
    /// string.try_push('b').unwrap();
    /// let overflow = string.try_push('c');
    ///
    /// assert_eq!(&string[..], "ab");
    /// assert_eq!(overflow.unwrap_err().element(), 'c');
    /// ```
    pub fn try_push(&mut self, c: char) -> Result<(), CapacityError<char>> {
        let len = self.len();
        unsafe {
            match encode_utf8(c, &mut self.raw_mut_bytes()[len..]) {
                Ok(n) => {
                    self.set_len(len + n);
                    Ok(())
                }
                Err(_) => Err(CapacityError::new(c)),
            }
        }
    }

    /// Adds the given string slice to the end of the string.
    ///
    /// Returns `Ok` if the push succeeds.
    ///
    /// ***Panics*** if the backing array is not large enough to fit the string.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 2]>::new();
    ///
    /// string.push_str("a");
    /// string.push_str("d");
    ///
    /// assert_eq!(&string[..], "ad");
    /// ```
    pub fn push_str(&mut self, s: &str) {
        self.try_push_str(s).unwrap()
    }

    /// Adds the given string slice to the end of the string.
    ///
    /// Returns `Ok` if the push succeeds.
    ///
    /// **Errors** if the backing array is not large enough to fit the string.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 2]>::new();
    ///
    /// string.try_push_str("a").unwrap();
    /// let overflow1 = string.try_push_str("bc");
    /// string.try_push_str("d").unwrap();
    /// let overflow2 = string.try_push_str("ef");
    ///
    /// assert_eq!(&string[..], "ad");
    /// assert_eq!(overflow1.unwrap_err().element(), "bc");
    /// assert_eq!(overflow2.unwrap_err().element(), "ef");
    /// ```
    pub fn try_push_str<'a>(&mut self, s: &'a str) -> Result<(), CapacityError<&'a str>> {
        if s.len() > self.capacity() - self.len() {
            return Err(CapacityError::new(s));
        }
        unsafe {
            let dst = self.xs.as_mut_ptr().offset(self.len() as isize);
            let src = s.as_ptr();
            ptr::copy_nonoverlapping(src, dst, s.len());
            let newl = self.len() + s.len();
            self.set_len(newl);
        }
        Ok(())
    }

    /// Make the string empty.
    pub fn clear(&mut self) {
        unsafe {
            self.set_len(0);
        }
    }

    /// Set the strings's length.
    ///
    /// May panic if `length` is greater than the capacity.
    ///
    /// This function is `unsafe` because it changes the notion of the
    /// number of “valid” bytes in the string. Use with care.
    #[inline]
    pub unsafe fn set_len(&mut self, length: usize) {
        debug_assert!(length <= self.capacity());
        self.len = Index::from(length);
    }

    /// Return a string slice of the whole `ArrayString`.
    pub fn as_str(&self) -> &str {
        self
    }

    /// Return a mutable slice of the whole string's buffer
    unsafe fn raw_mut_bytes(&mut self) -> &mut [u8] {
        slice::from_raw_parts_mut(self.xs.as_mut_ptr(), self.capacity())
    }
}

impl<A: Array<Item=u8>> Deref for ArrayString<A> {
    type Target = str;
    #[inline]
    fn deref(&self) -> &str {
        unsafe {
            let sl = slice::from_raw_parts(self.xs.as_ptr(), self.len.to_usize());
            str::from_utf8_unchecked(sl)
        }
    }
}

impl<A: Array<Item=u8>> DerefMut for ArrayString<A> {
    #[inline]
    fn deref_mut(&mut self) -> &mut str {
        unsafe {
            let sl = slice::from_raw_parts_mut(self.xs.as_mut_ptr(), self.len.to_usize());
            // FIXME: Nothing but transmute to do this right now
            mem::transmute(sl)
        }
    }
}

impl<A: Array<Item=u8>> PartialEq for ArrayString<A> {
    fn eq(&self, rhs: &Self) -> bool {
        **self == **rhs
    }
}

impl<A: Array<Item=u8>> PartialEq<str> for ArrayString<A> {
    fn eq(&self, rhs: &str) -> bool {
        &**self == rhs
    }
}

impl<A: Array<Item=u8>> PartialEq<ArrayString<A>> for str {
    fn eq(&self, rhs: &ArrayString<A>) -> bool {
        self == &**rhs
    }
}

impl<A: Array<Item=u8>> Eq for ArrayString<A> { }

impl<A: Array<Item=u8>> Hash for ArrayString<A> {
    fn hash<H: Hasher>(&self, h: &mut H) {
        (**self).hash(h)
    }
}

impl<A: Array<Item=u8>> Borrow<str> for ArrayString<A> {
    fn borrow(&self) -> &str { self }
}

impl<A: Array<Item=u8>> AsRef<str> for ArrayString<A> {
    fn as_ref(&self) -> &str { self }
}

impl<A: Array<Item=u8>> fmt::Debug for ArrayString<A> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { (**self).fmt(f) }
}

impl<A: Array<Item=u8>> fmt::Display for ArrayString<A> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { (**self).fmt(f) }
}

/// `Write` appends written data to the end of the string.
impl<A: Array<Item=u8>> fmt::Write for ArrayString<A> {
    fn write_char(&mut self, c: char) -> fmt::Result {
        self.try_push(c).map_err(|_| fmt::Error)
    }

    fn write_str(&mut self, s: &str) -> fmt::Result {
        self.try_push_str(s).map_err(|_| fmt::Error)
    }
}

impl<A: Array<Item=u8> + Copy> Clone for ArrayString<A> {
    fn clone(&self) -> ArrayString<A> {
        *self
    }
    fn clone_from(&mut self, rhs: &Self) {
        // guaranteed to fit due to types matching.
        self.clear();
        self.try_push_str(rhs).ok();
    }
}

impl<A: Array<Item=u8>> PartialOrd for ArrayString<A> {
    fn partial_cmp(&self, rhs: &Self) -> Option<cmp::Ordering> {
        (**self).partial_cmp(&**rhs)
    }
    fn lt(&self, rhs: &Self) -> bool { **self < **rhs }
    fn le(&self, rhs: &Self) -> bool { **self <= **rhs }
    fn gt(&self, rhs: &Self) -> bool { **self > **rhs }
    fn ge(&self, rhs: &Self) -> bool { **self >= **rhs }
}

impl<A: Array<Item=u8>> PartialOrd<str> for ArrayString<A> {
    fn partial_cmp(&self, rhs: &str) -> Option<cmp::Ordering> {
        (**self).partial_cmp(rhs)
    }
    fn lt(&self, rhs: &str) -> bool { &**self < rhs }
    fn le(&self, rhs: &str) -> bool { &**self <= rhs }
    fn gt(&self, rhs: &str) -> bool { &**self > rhs }
    fn ge(&self, rhs: &str) -> bool { &**self >= rhs }
}

impl<A: Array<Item=u8>> PartialOrd<ArrayString<A>> for str {
    fn partial_cmp(&self, rhs: &ArrayString<A>) -> Option<cmp::Ordering> {
        self.partial_cmp(&**rhs)
    }
    fn lt(&self, rhs: &ArrayString<A>) -> bool { self < &**rhs }
    fn le(&self, rhs: &ArrayString<A>) -> bool { self <= &**rhs }
    fn gt(&self, rhs: &ArrayString<A>) -> bool { self > &**rhs }
    fn ge(&self, rhs: &ArrayString<A>) -> bool { self >= &**rhs }
}

impl<A: Array<Item=u8>> Ord for ArrayString<A> {
    fn cmp(&self, rhs: &Self) -> cmp::Ordering {
        (**self).cmp(&**rhs)
    }
}

#[cfg(feature="serde-1")]
/// Requires crate feature `"serde-1"`
impl<A: Array<Item=u8>> Serialize for ArrayString<A> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer
    {
        serializer.serialize_str(&*self)
    }
}

#[cfg(feature="serde-1")]
/// Requires crate feature `"serde-1"`
impl<'de, A: Array<Item=u8>> Deserialize<'de> for ArrayString<A> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where D: Deserializer<'de>
    {
        use serde::de::{self, Visitor};
        use std::marker::PhantomData;

        struct ArrayStringVisitor<A: Array<Item=u8>>(PhantomData<A>);

        impl<'de, A: Array<Item=u8>> Visitor<'de> for ArrayStringVisitor<A> {
            type Value = ArrayString<A>;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                write!(formatter, "a string no more than {} bytes long", A::capacity())
            }

            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
                where E: de::Error,
            {
                ArrayString::from(v).map_err(|_| E::invalid_length(v.len(), &self))
            }

            fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
                where E: de::Error,
            {
                let s = try!(str::from_utf8(v).map_err(|_| E::invalid_value(de::Unexpected::Bytes(v), &self)));

                ArrayString::from(s).map_err(|_| E::invalid_length(s.len(), &self))
            }
        }

        deserializer.deserialize_str(ArrayStringVisitor::<A>(PhantomData))
    }
}
<commit_msg>DOC: Update ArrayString doc for push, push_str after 0.4 changes.<commit_after>use std::borrow::Borrow;
use std::cmp;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::mem;
use std::ptr;
use std::ops::{Deref, DerefMut};
use std::str;
use std::str::Utf8Error;
use std::slice;

use array::{Array, ArrayExt};
use array::Index;
use CapacityError;
use errors::PubCrateNew;
use odds::char::encode_utf8;

#[cfg(feature="serde-1")]
use serde::{Serialize, Deserialize, Serializer, Deserializer};

/// A string with a fixed capacity.
///
/// The `ArrayString` is a string backed by a fixed size array. It keeps track
/// of its length.
///
/// The string is a contiguous value that you can store directly on the stack
/// if needed.
#[derive(Copy)]
pub struct ArrayString<A: Array<Item=u8>> {
    xs: A,
    len: A::Index,
}

impl<A: Array<Item=u8>> ArrayString<A> {
    /// Create a new empty `ArrayString`.
    ///
    /// Capacity is inferred from the type parameter.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 16]>::new();
    /// string.push_str("foo");
    /// assert_eq!(&string[..], "foo");
    /// assert_eq!(string.capacity(), 16);
    /// ```
    pub fn new() -> ArrayString<A> {
        unsafe {
            ArrayString {
                xs: ::new_array(),
                len: Index::from(0),
            }
        }
    }

    /// Create a new `ArrayString` from a `str`.
    ///
    /// Capacity is inferred from the type parameter.
    ///
    /// **Errors** if the backing array is not large enough to fit the string.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 3]>::from("foo").unwrap();
    /// assert_eq!(&string[..], "foo");
    /// assert_eq!(string.len(), 3);
    /// assert_eq!(string.capacity(), 3);
    /// ```
    pub fn from(s: &str) -> Result<Self, CapacityError<&str>> {
        let mut arraystr = Self::new();
        arraystr.try_push_str(s)?;
        Ok(arraystr)
    }

    /// Create a new `ArrayString` from a byte string literal.
    ///
    /// **Errors** if the byte string literal is not valid UTF-8.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let string = ArrayString::from_byte_string(b"hello world").unwrap();
    /// ```
    pub fn from_byte_string(b: &A) -> Result<Self, Utf8Error> {
        let mut arraystr = Self::new();
        let s = try!(str::from_utf8(b.as_slice()));
        let _result = arraystr.try_push_str(s);
        debug_assert!(_result.is_ok());
        Ok(arraystr)
    }

    /// Return the capacity of the `ArrayString`.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let string = ArrayString::<[_; 3]>::new();
    /// assert_eq!(string.capacity(), 3);
    /// ```
    #[inline]
    pub fn capacity(&self) -> usize { A::capacity() }

    /// Return if the `ArrayString` is completely filled.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 1]>::new();
    /// assert!(!string.is_full());
    /// string.push_str("A");
    /// assert!(string.is_full());
    /// ```
    pub fn is_full(&self) -> bool { self.len() == self.capacity() }

    /// Adds the given char to the end of the string.
    ///
    /// ***Panics*** if the backing array is not large enough to fit the additional char.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 2]>::new();
    ///
    /// string.push('a');
    /// string.push('b');
    ///
    /// assert_eq!(&string[..], "ab");
    /// ```
    pub fn push(&mut self, c: char) {
        self.try_push(c).unwrap();
    }

    /// Adds the given char to the end of the string.
    ///
    /// Returns `Ok` if the push succeeds.
    ///
    /// **Errors** if the backing array is not large enough to fit the additional char.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 2]>::new();
    ///
    /// string.try_push('a').unwrap();
    /// string.try_push('b').unwrap();
    /// let overflow = string.try_push('c');
    ///
    /// assert_eq!(&string[..], "ab");
    /// assert_eq!(overflow.unwrap_err().element(), 'c');
    /// ```
    pub fn try_push(&mut self, c: char) -> Result<(), CapacityError<char>> {
        let len = self.len();
        unsafe {
            match encode_utf8(c, &mut self.raw_mut_bytes()[len..]) {
                Ok(n) => {
                    self.set_len(len + n);
                    Ok(())
                }
                Err(_) => Err(CapacityError::new(c)),
            }
        }
    }

    /// Adds the given string slice to the end of the string.
    ///
    /// ***Panics*** if the backing array is not large enough to fit the string.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 2]>::new();
    ///
    /// string.push_str("a");
    /// string.push_str("d");
    ///
    /// assert_eq!(&string[..], "ad");
    /// ```
    pub fn push_str(&mut self, s: &str) {
        self.try_push_str(s).unwrap()
    }

    /// Adds the given string slice to the end of the string.
    ///
    /// Returns `Ok` if the push succeeds.
    ///
    /// **Errors** if the backing array is not large enough to fit the string.
    ///
    /// ```
    /// use arrayvec::ArrayString;
    ///
    /// let mut string = ArrayString::<[_; 2]>::new();
    ///
    /// string.try_push_str("a").unwrap();
    /// let overflow1 = string.try_push_str("bc");
    /// string.try_push_str("d").unwrap();
    /// let overflow2 = string.try_push_str("ef");
    ///
    /// assert_eq!(&string[..], "ad");
    /// assert_eq!(overflow1.unwrap_err().element(), "bc");
    /// assert_eq!(overflow2.unwrap_err().element(), "ef");
    /// ```
    pub fn try_push_str<'a>(&mut self, s: &'a str) -> Result<(), CapacityError<&'a str>> {
        if s.len() > self.capacity() - self.len() {
            return Err(CapacityError::new(s));
        }
        unsafe {
            let dst = self.xs.as_mut_ptr().offset(self.len() as isize);
            let src = s.as_ptr();
            ptr::copy_nonoverlapping(src, dst, s.len());
            let newl = self.len() + s.len();
            self.set_len(newl);
        }
        Ok(())
    }

    /// Make the string empty.
    pub fn clear(&mut self) {
        unsafe {
            self.set_len(0);
        }
    }

    /// Set the strings's length.
    ///
    /// May panic if `length` is greater than the capacity.
    ///
    /// This function is `unsafe` because it changes the notion of the
    /// number of “valid” bytes in the string. Use with care.
    #[inline]
    pub unsafe fn set_len(&mut self, length: usize) {
        debug_assert!(length <= self.capacity());
        self.len = Index::from(length);
    }

    /// Return a string slice of the whole `ArrayString`.
    pub fn as_str(&self) -> &str {
        self
    }

    /// Return a mutable slice of the whole string's buffer
    unsafe fn raw_mut_bytes(&mut self) -> &mut [u8] {
        slice::from_raw_parts_mut(self.xs.as_mut_ptr(), self.capacity())
    }
}

impl<A: Array<Item=u8>> Deref for ArrayString<A> {
    type Target = str;
    #[inline]
    fn deref(&self) -> &str {
        unsafe {
            let sl = slice::from_raw_parts(self.xs.as_ptr(), self.len.to_usize());
            str::from_utf8_unchecked(sl)
        }
    }
}

impl<A: Array<Item=u8>> DerefMut for ArrayString<A> {
    #[inline]
    fn deref_mut(&mut self) -> &mut str {
        unsafe {
            let sl = slice::from_raw_parts_mut(self.xs.as_mut_ptr(), self.len.to_usize());
            // FIXME: Nothing but transmute to do this right now
            mem::transmute(sl)
        }
    }
}

impl<A: Array<Item=u8>> PartialEq for ArrayString<A> {
    fn eq(&self, rhs: &Self) -> bool {
        **self == **rhs
    }
}

impl<A: Array<Item=u8>> PartialEq<str> for ArrayString<A> {
    fn eq(&self, rhs: &str) -> bool {
        &**self == rhs
    }
}

impl<A: Array<Item=u8>> PartialEq<ArrayString<A>> for str {
    fn eq(&self, rhs: &ArrayString<A>) -> bool {
        self == &**rhs
    }
}

impl<A: Array<Item=u8>> Eq for ArrayString<A> { }

impl<A: Array<Item=u8>> Hash for ArrayString<A> {
    fn hash<H: Hasher>(&self, h: &mut H) {
        (**self).hash(h)
    }
}

impl<A: Array<Item=u8>> Borrow<str> for ArrayString<A> {
    fn borrow(&self) -> &str { self }
}

impl<A: Array<Item=u8>> AsRef<str> for ArrayString<A> {
    fn as_ref(&self) -> &str { self }
}

impl<A: Array<Item=u8>> fmt::Debug for ArrayString<A> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { (**self).fmt(f) }
}

impl<A: Array<Item=u8>> fmt::Display for ArrayString<A> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { (**self).fmt(f) }
}

/// `Write` appends written data to the end of the string.
impl<A: Array<Item=u8>> fmt::Write for ArrayString<A> {
    fn write_char(&mut self, c: char) -> fmt::Result {
        self.try_push(c).map_err(|_| fmt::Error)
    }

    fn write_str(&mut self, s: &str) -> fmt::Result {
        self.try_push_str(s).map_err(|_| fmt::Error)
    }
}

impl<A: Array<Item=u8> + Copy> Clone for ArrayString<A> {
    fn clone(&self) -> ArrayString<A> {
        *self
    }
    fn clone_from(&mut self, rhs: &Self) {
        // guaranteed to fit due to types matching.
        self.clear();
        self.try_push_str(rhs).ok();
    }
}

impl<A: Array<Item=u8>> PartialOrd for ArrayString<A> {
    fn partial_cmp(&self, rhs: &Self) -> Option<cmp::Ordering> {
        (**self).partial_cmp(&**rhs)
    }
    fn lt(&self, rhs: &Self) -> bool { **self < **rhs }
    fn le(&self, rhs: &Self) -> bool { **self <= **rhs }
    fn gt(&self, rhs: &Self) -> bool { **self > **rhs }
    fn ge(&self, rhs: &Self) -> bool { **self >= **rhs }
}

impl<A: Array<Item=u8>> PartialOrd<str> for ArrayString<A> {
    fn partial_cmp(&self, rhs: &str) -> Option<cmp::Ordering> {
        (**self).partial_cmp(rhs)
    }
    fn lt(&self, rhs: &str) -> bool { &**self < rhs }
    fn le(&self, rhs: &str) -> bool { &**self <= rhs }
    fn gt(&self, rhs: &str) -> bool { &**self > rhs }
    fn ge(&self, rhs: &str) -> bool { &**self >= rhs }
}

impl<A: Array<Item=u8>> PartialOrd<ArrayString<A>> for str {
    fn partial_cmp(&self, rhs: &ArrayString<A>) -> Option<cmp::Ordering> {
        self.partial_cmp(&**rhs)
    }
    fn lt(&self, rhs: &ArrayString<A>) -> bool { self < &**rhs }
    fn le(&self, rhs: &ArrayString<A>) -> bool { self <= &**rhs }
    fn gt(&self, rhs: &ArrayString<A>) -> bool { self > &**rhs }
    fn ge(&self, rhs: &ArrayString<A>) -> bool { self >= &**rhs }
}

impl<A: Array<Item=u8>> Ord for ArrayString<A> {
    fn cmp(&self, rhs: &Self) -> cmp::Ordering {
        (**self).cmp(&**rhs)
    }
}

#[cfg(feature="serde-1")]
/// Requires crate feature `"serde-1"`
impl<A: Array<Item=u8>> Serialize for ArrayString<A> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
        where S: Serializer
    {
        serializer.serialize_str(&*self)
    }
}

#[cfg(feature="serde-1")]
/// Requires crate feature `"serde-1"`
impl<'de, A: Array<Item=u8>> Deserialize<'de> for ArrayString<A> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
        where D: Deserializer<'de>
    {
        use serde::de::{self, Visitor};
        use std::marker::PhantomData;

        struct ArrayStringVisitor<A: Array<Item=u8>>(PhantomData<A>);

        impl<'de, A: Array<Item=u8>> Visitor<'de> for ArrayStringVisitor<A> {
            type Value = ArrayString<A>;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                write!(formatter, "a string no more than {} bytes long", A::capacity())
            }

            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
                where E: de::Error,
            {
                ArrayString::from(v).map_err(|_| E::invalid_length(v.len(), &self))
            }

            fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
                where E: de::Error,
            {
                let s = try!(str::from_utf8(v).map_err(|_| E::invalid_value(de::Unexpected::Bytes(v), &self)));

                ArrayString::from(s).map_err(|_| E::invalid_length(s.len(), &self))
            }
        }

        deserializer.deserialize_str(ArrayStringVisitor::<A>(PhantomData))
    }
}
<|endoftext|>


<commit_before>// library ~ (core/bundler file)

// Copyright (C) ~ Alex Lyon <arcterus@mail.com>
// Copyright (C) ~ Roy Ivy III <rivy.dev@gmail.com>; MIT license

// * feature-gated external crates
#[cfg(all(feature = "lazy_static", target_os = "linux"))]
extern crate lazy_static;
#[cfg(feature = "nix")]
extern crate nix;
#[cfg(feature = "platform-info")]
extern crate platform_info;

// * feature-gated external crates (re-shared as public internal modules)
#[cfg(feature = "libc")]
pub extern crate libc;
#[cfg(feature = "winapi")]
pub extern crate winapi;

//## internal modules

mod features; // feature-gated code modules
mod macros; // crate macros (macro_rules-type; exported to `crate::...`)
mod mods; // core cross-platform modules
mod parser; // string parsing modules

// * cross-platform modules
pub use crate::mods::backup_control;
pub use crate::mods::coreopts;
pub use crate::mods::os;
pub use crate::mods::panic;
pub use crate::mods::ranges;

// * string parsing modules
pub use crate::parser::parse_size;
pub use crate::parser::parse_time;

// * feature-gated modules
#[cfg(feature = "encoding")]
pub use crate::features::encoding;
#[cfg(feature = "fs")]
pub use crate::features::fs;
#[cfg(feature = "fsext")]
pub use crate::features::fsext;
#[cfg(feature = "ringbuffer")]
pub use crate::features::ringbuffer;
#[cfg(feature = "zero-copy")]
pub use crate::features::zero_copy;

// * (platform-specific) feature-gated modules
// ** non-windows
#[cfg(all(not(windows), feature = "mode"))]
pub use crate::features::mode;
// ** unix-only
#[cfg(all(unix, feature = "entries"))]
pub use crate::features::entries;
#[cfg(all(unix, feature = "perms"))]
pub use crate::features::perms;
#[cfg(all(unix, feature = "process"))]
pub use crate::features::process;
#[cfg(all(unix, not(target_os = "fuchsia"), feature = "signals"))]
pub use crate::features::signals;
#[cfg(all(
    unix,
    not(target_os = "fuchsia"),
    not(target_env = "musl"),
    feature = "utmpx"
))]
pub use crate::features::utmpx;
// ** windows-only
#[cfg(all(windows, feature = "wide"))]
pub use crate::features::wide;

//## core functions

use std::ffi::OsString;

pub enum InvalidEncodingHandling {
    Ignore,
    ConvertLossy,
    Panic,
}

#[must_use]
pub enum ConversionResult {
    Complete(Vec<String>),
    Lossy(Vec<String>),
}

impl ConversionResult {
    pub fn accept_any(self) -> Vec<String> {
        match self {
            Self::Complete(result) => result,
            Self::Lossy(result) => result,
        }
    }

    pub fn expect_lossy(self, msg: &str) -> Vec<String> {
        match self {
            Self::Lossy(result) => result,
            Self::Complete(_) => {
                panic!("{}", msg);
            }
        }
    }

    pub fn expect_complete(self, msg: &str) -> Vec<String> {
        match self {
            Self::Complete(result) => result,
            Self::Lossy(_) => {
                panic!("{}", msg);
            }
        }
    }
}

pub trait Args: Iterator<Item = OsString> + Sized {
    /// Converts each iterator item to a String and collects these into a vector
    /// On invalid encoding, the result will depend on the argument. This method allows to either drop entries with illegal encoding
    /// completely (```InvalidEncodingHandling::Ignore```), convert them using lossy-conversion (```InvalidEncodingHandling::Lossy```) which will
    /// result in strange strings or can chosen to panic (```InvalidEncodingHandling::Panic```).
    /// # Arguments
    /// * `handling` - This switch allows to switch the behavior, when invalid encoding is encountered
    /// # Panics
    /// * Occurs, when invalid encoding is encountered and handling is set to ```InvalidEncodingHandling::Panic```
    fn collect_str(self, handling: InvalidEncodingHandling) -> ConversionResult {
        let mut full_conversion = true;
        let result_vector: Vec<String> = self
            .map(|s| match s.into_string() {
                Ok(string) => Ok(string),
                Err(s_ret) => {
                    full_conversion = false;
                    let lossy_conversion = s_ret.to_string_lossy();
                    eprintln!(
                        "Input with broken encoding occurred! (s = '{}') ",
                        &lossy_conversion
                    );
                    match handling {
                        InvalidEncodingHandling::Ignore => Err(String::new()),
                        InvalidEncodingHandling::ConvertLossy => Err(lossy_conversion.to_string()),
                        InvalidEncodingHandling::Panic => {
                            panic!("Broken encoding found but caller cannot handle it")
                        }
                    }
                }
            })
            .filter(|s| match handling {
                InvalidEncodingHandling::Ignore => s.is_ok(),
                _ => true,
            })
            .map(|s| match s {
                Ok(v) => v,
                Err(e) => e,
            })
            .collect();

        if full_conversion {
            ConversionResult::Complete(result_vector)
        } else {
            ConversionResult::Lossy(result_vector)
        }
    }

    /// convenience function for a more slim interface
    fn collect_str_lossy(self) -> ConversionResult {
        self.collect_str(InvalidEncodingHandling::ConvertLossy)
    }
}

impl<T: Iterator<Item = OsString> + Sized> Args for T {}

// args() ...
pub fn args() -> impl Iterator<Item = String> {
    wild::args()
}

pub fn args_os() -> impl Iterator<Item = OsString> {
    wild::args_os()
}

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

    fn make_os_vec(os_str: &OsStr) -> Vec<OsString> {
        vec![
            OsString::from("test"),
            OsString::from("สวัสดี"),
            os_str.to_os_string(),
        ]
    }

    fn collect_os_str(vec: Vec<OsString>, handling: InvalidEncodingHandling) -> ConversionResult {
        vec.into_iter().collect_str(handling)
    }

    #[cfg(any(unix, target_os = "redox"))]
    fn test_invalid_utf8_args_lossy(os_str: &OsStr) {
        //assert our string is invalid utf8
        assert!(os_str.to_os_string().into_string().is_err());
        let test_vec = make_os_vec(os_str);
        let collected_to_str =
            collect_os_str(test_vec.clone(), InvalidEncodingHandling::ConvertLossy)
                .expect_lossy("Lossy conversion expected in this test: bad encoding entries should be converted as good as possible");
        //conservation of length - when accepting lossy conversion no arguments may be dropped
        assert_eq!(collected_to_str.len(), test_vec.len());
        //first indices identical
        for index in 0..2 {
            assert_eq!(
                collected_to_str.get(index).unwrap(),
                test_vec.get(index).unwrap().to_str().unwrap()
            );
        }
        //lossy conversion for string with illegal encoding is done
        assert_eq!(
            *collected_to_str.get(2).unwrap(),
            os_str.to_os_string().to_string_lossy()
        );
    }

    #[cfg(any(unix, target_os = "redox"))]
    fn test_invalid_utf8_args_ignore(os_str: &OsStr) {
        //assert our string is invalid utf8
        assert!(os_str.to_os_string().into_string().is_err());
        let test_vec = make_os_vec(os_str);
        let collected_to_str = collect_os_str(test_vec.clone(), InvalidEncodingHandling::Ignore)
            .expect_lossy(
                "Lossy conversion expected in this test: bad encoding entries should be filtered",
            );
        //assert that the broken entry is filtered out
        assert_eq!(collected_to_str.len(), test_vec.len() - 1);
        //assert that the unbroken indices are converted as expected
        for index in 0..2 {
            assert_eq!(
                collected_to_str.get(index).unwrap(),
                test_vec.get(index).unwrap().to_str().unwrap()
            );
        }
    }

    #[test]
    fn valid_utf8_encoding_args() {
        //create a vector containing only correct encoding
        let test_vec = make_os_vec(&OsString::from("test2"));
        //expect complete conversion without losses, even when lossy conversion is accepted
        let _ = collect_os_str(test_vec, InvalidEncodingHandling::ConvertLossy)
            .expect_complete("Lossy conversion not expected in this test");
    }

    #[cfg(any(unix, target_os = "redox"))]
    #[test]
    fn invalid_utf8_args_unix() {
        use std::os::unix::ffi::OsStrExt;

        let source = [0x66, 0x6f, 0x80, 0x6f];
        let os_str = OsStr::from_bytes(&source[..]);
        test_invalid_utf8_args_lossy(os_str);
        test_invalid_utf8_args_ignore(os_str);
    }
}
<commit_msg>docs/spell ~ (uucore) add spelling exceptions<commit_after>// library ~ (core/bundler file)

// Copyright (C) ~ Alex Lyon <arcterus@mail.com>
// Copyright (C) ~ Roy Ivy III <rivy.dev@gmail.com>; MIT license

// * feature-gated external crates
#[cfg(all(feature = "lazy_static", target_os = "linux"))]
extern crate lazy_static;
#[cfg(feature = "nix")]
extern crate nix;
#[cfg(feature = "platform-info")]
extern crate platform_info;

// * feature-gated external crates (re-shared as public internal modules)
#[cfg(feature = "libc")]
pub extern crate libc;
#[cfg(feature = "winapi")]
pub extern crate winapi;

//## internal modules

mod features; // feature-gated code modules
mod macros; // crate macros (macro_rules-type; exported to `crate::...`)
mod mods; // core cross-platform modules
mod parser; // string parsing modules

// * cross-platform modules
pub use crate::mods::backup_control;
pub use crate::mods::coreopts;
pub use crate::mods::os;
pub use crate::mods::panic;
pub use crate::mods::ranges;

// * string parsing modules
pub use crate::parser::parse_size;
pub use crate::parser::parse_time;

// * feature-gated modules
#[cfg(feature = "encoding")]
pub use crate::features::encoding;
#[cfg(feature = "fs")]
pub use crate::features::fs;
#[cfg(feature = "fsext")]
pub use crate::features::fsext;
#[cfg(feature = "ringbuffer")]
pub use crate::features::ringbuffer;
#[cfg(feature = "zero-copy")]
pub use crate::features::zero_copy;

// * (platform-specific) feature-gated modules
// ** non-windows
#[cfg(all(not(windows), feature = "mode"))]
pub use crate::features::mode;
// ** unix-only
#[cfg(all(unix, feature = "entries"))]
pub use crate::features::entries;
#[cfg(all(unix, feature = "perms"))]
pub use crate::features::perms;
#[cfg(all(unix, feature = "process"))]
pub use crate::features::process;
#[cfg(all(unix, not(target_os = "fuchsia"), feature = "signals"))]
pub use crate::features::signals;
#[cfg(all(
    unix,
    not(target_os = "fuchsia"),
    not(target_env = "musl"),
    feature = "utmpx"
))]
pub use crate::features::utmpx;
// ** windows-only
#[cfg(all(windows, feature = "wide"))]
pub use crate::features::wide;

//## core functions

use std::ffi::OsString;

pub enum InvalidEncodingHandling {
    Ignore,
    ConvertLossy,
    Panic,
}

#[must_use]
pub enum ConversionResult {
    Complete(Vec<String>),
    Lossy(Vec<String>),
}

impl ConversionResult {
    pub fn accept_any(self) -> Vec<String> {
        match self {
            Self::Complete(result) => result,
            Self::Lossy(result) => result,
        }
    }

    pub fn expect_lossy(self, msg: &str) -> Vec<String> {
        match self {
            Self::Lossy(result) => result,
            Self::Complete(_) => {
                panic!("{}", msg);
            }
        }
    }

    pub fn expect_complete(self, msg: &str) -> Vec<String> {
        match self {
            Self::Complete(result) => result,
            Self::Lossy(_) => {
                panic!("{}", msg);
            }
        }
    }
}

pub trait Args: Iterator<Item = OsString> + Sized {
    /// Converts each iterator item to a String and collects these into a vector
    /// On invalid encoding, the result will depend on the argument. This method allows to either drop entries with illegal encoding
    /// completely (```InvalidEncodingHandling::Ignore```), convert them using lossy-conversion (```InvalidEncodingHandling::Lossy```) which will
    /// result in strange strings or can chosen to panic (```InvalidEncodingHandling::Panic```).
    /// # Arguments
    /// * `handling` - This switch allows to switch the behavior, when invalid encoding is encountered
    /// # Panics
    /// * Occurs, when invalid encoding is encountered and handling is set to ```InvalidEncodingHandling::Panic```
    fn collect_str(self, handling: InvalidEncodingHandling) -> ConversionResult {
        let mut full_conversion = true;
        let result_vector: Vec<String> = self
            .map(|s| match s.into_string() {
                Ok(string) => Ok(string),
                Err(s_ret) => {
                    full_conversion = false;
                    let lossy_conversion = s_ret.to_string_lossy();
                    eprintln!(
                        "Input with broken encoding occurred! (s = '{}') ",
                        &lossy_conversion
                    );
                    match handling {
                        InvalidEncodingHandling::Ignore => Err(String::new()),
                        InvalidEncodingHandling::ConvertLossy => Err(lossy_conversion.to_string()),
                        InvalidEncodingHandling::Panic => {
                            panic!("Broken encoding found but caller cannot handle it")
                        }
                    }
                }
            })
            .filter(|s| match handling {
                InvalidEncodingHandling::Ignore => s.is_ok(),
                _ => true,
            })
            .map(|s| match s {
                Ok(v) => v,
                Err(e) => e,
            })
            .collect();

        if full_conversion {
            ConversionResult::Complete(result_vector)
        } else {
            ConversionResult::Lossy(result_vector)
        }
    }

    /// convenience function for a more slim interface
    fn collect_str_lossy(self) -> ConversionResult {
        self.collect_str(InvalidEncodingHandling::ConvertLossy)
    }
}

impl<T: Iterator<Item = OsString> + Sized> Args for T {}

// args() ...
pub fn args() -> impl Iterator<Item = String> {
    wild::args()
}

pub fn args_os() -> impl Iterator<Item = OsString> {
    wild::args_os()
}

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

    fn make_os_vec(os_str: &OsStr) -> Vec<OsString> {
        vec![
            OsString::from("test"),
            OsString::from("สวัสดี"), // spell-checker:disable-line
            os_str.to_os_string(),
        ]
    }

    fn collect_os_str(vec: Vec<OsString>, handling: InvalidEncodingHandling) -> ConversionResult {
        vec.into_iter().collect_str(handling)
    }

    #[cfg(any(unix, target_os = "redox"))]
    fn test_invalid_utf8_args_lossy(os_str: &OsStr) {
        //assert our string is invalid utf8
        assert!(os_str.to_os_string().into_string().is_err());
        let test_vec = make_os_vec(os_str);
        let collected_to_str =
            collect_os_str(test_vec.clone(), InvalidEncodingHandling::ConvertLossy)
                .expect_lossy("Lossy conversion expected in this test: bad encoding entries should be converted as good as possible");
        //conservation of length - when accepting lossy conversion no arguments may be dropped
        assert_eq!(collected_to_str.len(), test_vec.len());
        //first indices identical
        for index in 0..2 {
            assert_eq!(
                collected_to_str.get(index).unwrap(),
                test_vec.get(index).unwrap().to_str().unwrap()
            );
        }
        //lossy conversion for string with illegal encoding is done
        assert_eq!(
            *collected_to_str.get(2).unwrap(),
            os_str.to_os_string().to_string_lossy()
        );
    }

    #[cfg(any(unix, target_os = "redox"))]
    fn test_invalid_utf8_args_ignore(os_str: &OsStr) {
        //assert our string is invalid utf8
        assert!(os_str.to_os_string().into_string().is_err());
        let test_vec = make_os_vec(os_str);
        let collected_to_str = collect_os_str(test_vec.clone(), InvalidEncodingHandling::Ignore)
            .expect_lossy(
                "Lossy conversion expected in this test: bad encoding entries should be filtered",
            );
        //assert that the broken entry is filtered out
        assert_eq!(collected_to_str.len(), test_vec.len() - 1);
        //assert that the unbroken indices are converted as expected
        for index in 0..2 {
            assert_eq!(
                collected_to_str.get(index).unwrap(),
                test_vec.get(index).unwrap().to_str().unwrap()
            );
        }
    }

    #[test]
    fn valid_utf8_encoding_args() {
        //create a vector containing only correct encoding
        let test_vec = make_os_vec(&OsString::from("test2"));
        //expect complete conversion without losses, even when lossy conversion is accepted
        let _ = collect_os_str(test_vec, InvalidEncodingHandling::ConvertLossy)
            .expect_complete("Lossy conversion not expected in this test");
    }

    #[cfg(any(unix, target_os = "redox"))]
    #[test]
    fn invalid_utf8_args_unix() {
        use std::os::unix::ffi::OsStrExt;

        let source = [0x66, 0x6f, 0x80, 0x6f];
        let os_str = OsStr::from_bytes(&source[..]);
        test_invalid_utf8_args_lossy(os_str);
        test_invalid_utf8_args_ignore(os_str);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;

extern crate arguments;
extern crate random;
extern crate sql;
extern crate sqlite;
extern crate streamer;
extern crate term;
extern crate time;

use log::LogLevel;
use std::error::Error;
use streamer::{Config, Result, System};

const USAGE: &'static str = "
Usage: streamer [options]

Options:
    --config <path>          Configuration file (required).
    --length <time>          Time span to simulate in seconds [default: 10].
    --output <path>          Output file for power and temperature profiles.

    --verbose                Display progress information.
    --help                   Display this message.
";

macro_rules! raise(
    ($message:expr) => (return Err(Box::new(::streamer::ErrorString($message.to_string()))));
);

macro_rules! ok(
    ($result:expr) => (match $result {
        Ok(result) => result,
        Err(error) => return Err(Box::new(error)),
    });
);

macro_rules! some(
    ($option:expr, $($arg:tt)*) => (match $option {
        Some(value) => value,
        _ => raise!($($arg)*),
    });
);

mod logger;
mod output;

use output::Output;

fn main() {
    start().unwrap_or_else(|error| fail(&*error));
}

fn start() -> Result<()> {
    let arguments = ok!(arguments::parse(std::env::args()));

    if arguments.get::<bool>("help").unwrap_or(false) {
        help();
    }

    if arguments.get::<bool>("verbose").unwrap_or(false) {
        logger::setup(LogLevel::Info);
    } else {
        logger::setup(LogLevel::Warn);
    }

    let mut system = {
        let config = some!(arguments.get::<String>("config"), "a configuration file is required");
        let config = try!(Config::new(config));
        let seed = config.get::<i64>("seed").map(|&seed| seed as u64).unwrap_or(0);
        let seed = if seed > 0 { seed } else { 0x04020609 };
        let seed = [0x12345678 & seed, 0x87654321 & seed];
        try!(System::new(&config, &random::default().seed(seed)))
    };

    let length = arguments.get::<f64>("length").unwrap_or(10.0);
    if length <= 0.0 {
        raise!("the time span should be positive");
    }

    let mut output = try!(output::new(&system, arguments.get::<String>("output")));

    info!(target: "Streamer", "Simulating {} seconds with a time step of {} seconds...", length,
          system.time_step());
    let start = time::now();

    while let Some((event, power, temperature)) = system.next() {
        if event.time > length {
            break;
        }
        info!(target: "Streamer", "{} | {:2} queued", event,
              system.stats.arrived - system.stats.started);
        try!(output.next((event, power, temperature)));
    }

    info!(target: "Streamer", "Well done in {:.2} seconds.",
          (time::now() - start).num_milliseconds() as f64 / 1000.0);

    Ok(())
}

fn help() -> ! {
    println!("{}", USAGE.trim());
    std::process::exit(0);
}

fn fail(error: &Error) -> ! {
    use std::io::{stderr, Write};
    stderr().write_all(format!("Error: {}.\n", error).as_bytes()).unwrap();
    std::process::exit(1);
}
<commit_msg>Initialize the RNG based on time if seed is zero<commit_after>#[macro_use]
extern crate log;

extern crate arguments;
extern crate random;
extern crate sql;
extern crate sqlite;
extern crate streamer;
extern crate term;
extern crate time;

use log::LogLevel;
use std::error::Error;
use streamer::{Config, Result, System};

const USAGE: &'static str = "
Usage: streamer [options]

Options:
    --config <path>          Configuration file (required).
    --length <time>          Time span to simulate in seconds [default: 10].
    --output <path>          Output file for power and temperature profiles.

    --verbose                Display progress information.
    --help                   Display this message.
";

macro_rules! raise(
    ($message:expr) => (return Err(Box::new(::streamer::ErrorString($message.to_string()))));
);

macro_rules! ok(
    ($result:expr) => (match $result {
        Ok(result) => result,
        Err(error) => return Err(Box::new(error)),
    });
);

macro_rules! some(
    ($option:expr, $($arg:tt)*) => (match $option {
        Some(value) => value,
        _ => raise!($($arg)*),
    });
);

mod logger;
mod output;

use output::Output;

fn main() {
    start().unwrap_or_else(|error| fail(&*error));
}

fn start() -> Result<()> {
    let arguments = ok!(arguments::parse(std::env::args()));

    if arguments.get::<bool>("help").unwrap_or(false) {
        help();
    }

    if arguments.get::<bool>("verbose").unwrap_or(false) {
        logger::setup(LogLevel::Info);
    } else {
        logger::setup(LogLevel::Warn);
    }

    let mut system = {
        let config = some!(arguments.get::<String>("config"), "a configuration file is required");
        let config = try!(Config::new(config));
        let seed = config.get::<i64>("seed").map(|&seed| seed as u64).unwrap_or(0);
        let seed = if seed > 0 { seed } else { time::now().to_timespec().sec as u64 };
        let seed = [0x12345678 & seed, 0x87654321 & seed];
        try!(System::new(&config, &random::default().seed(seed)))
    };

    let length = arguments.get::<f64>("length").unwrap_or(10.0);
    if length <= 0.0 {
        raise!("the time span should be positive");
    }

    let mut output = try!(output::new(&system, arguments.get::<String>("output")));

    info!(target: "Streamer", "Simulating {} seconds with a time step of {} seconds...", length,
          system.time_step());
    let start = time::now();

    while let Some((event, power, temperature)) = system.next() {
        if event.time > length {
            break;
        }
        info!(target: "Streamer", "{} | {:2} queued", event,
              system.stats.arrived - system.stats.started);
        try!(output.next((event, power, temperature)));
    }

    info!(target: "Streamer", "Well done in {:.2} seconds.",
          (time::now() - start).num_milliseconds() as f64 / 1000.0);

    Ok(())
}

fn help() -> ! {
    println!("{}", USAGE.trim());
    std::process::exit(0);
}

fn fail(error: &Error) -> ! {
    use std::io::{stderr, Write};
    stderr().write_all(format!("Error: {}.\n", error).as_bytes()).unwrap();
    std::process::exit(1);
}
<|endoftext|>


<commit_before>use crate::error::{Error,
                   Result};
pub use crate::os::system::{uname,
                            Uname};
use std::{io,
          net::{IpAddr,
                Ipv4Addr,
                SocketAddr,
                ToSocketAddrs,
                UdpSocket}};

const UNSPECIFIED_SOCKET_ADDR: (Ipv4Addr, u16) = (Ipv4Addr::UNSPECIFIED, 0);

crate::env_config_socketaddr!(OutboundIpAddrLookupSocketAddr,
                              HAB_OUTBOUND_IP_ADDR_LOOKUP_SOCKET_ADDR,
                              // Use Google DNS as the default
                              8,
                              8,
                              8,
                              8,
                              53);

pub fn ip() -> Result<IpAddr> {
    let connect_addr = SocketAddr::from(OutboundIpAddrLookupSocketAddr::configured_value());
    ip_impl(connect_addr).map_err(Error::NoOutboundIpAddr)
}

fn ip_impl(connect_addr: impl ToSocketAddrs) -> io::Result<IpAddr> {
    let socket = UdpSocket::bind(UNSPECIFIED_SOCKET_ADDR)?;
    socket.connect(connect_addr)?;
    let addr = socket.local_addr()?;
    Ok(addr.ip())
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn ip_lookup() {
        assert!(ip_impl("an invalid socket addr").is_err());
        assert!(ip().is_ok());

        // We would prefer to use the `locked_env_var` macro here but currently, the `habitat_core`
        // crate does not depend on `habitat_common`. If this becomes an issue, some refactoring
        // will be needed.
        std::env::set_var(OutboundIpAddrLookupSocketAddr::ENVVAR, "127.0.0.1:80");
        // localhost should route back to localhost
        assert_eq!(ip().unwrap(), Ipv4Addr::LOCALHOST);
    }
}
<commit_msg>Add doc comment on IP lookup technique<commit_after>use crate::error::{Error,
                   Result};
pub use crate::os::system::{uname,
                            Uname};
use std::{io,
          net::{IpAddr,
                Ipv4Addr,
                SocketAddr,
                ToSocketAddrs,
                UdpSocket}};

const UNSPECIFIED_SOCKET_ADDR: (Ipv4Addr, u16) = (Ipv4Addr::UNSPECIFIED, 0);

crate::env_config_socketaddr!(OutboundIpAddrLookupSocketAddr,
                              HAB_OUTBOUND_IP_ADDR_LOOKUP_SOCKET_ADDR,
                              // Use Google DNS as the default
                              8,
                              8,
                              8,
                              8,
                              53);

/// The technique used to determine the outgoing IP address is documented [here][1].
///
/// "A connected UDP socket can also be used to determine the outgoing interface
/// that will be used to a particular destination. This is because of a side effect of the connect
/// function when applied to a UDP socket: The kernel chooses the local IP address (assuming the
/// process has not already called bind to explicitly assign this). This local IP address is chosen
/// by searching the routing table for the destination IP address, and then using the primary IP
/// address for the resulting interface." From "Unix Network Programming v1.3" chapter 8
/// section 14.
///
/// [1]: http://www.masterraghu.com/subjects/np/introduction/unix_network_programming_v1.3/ch08lev1sec14.html
pub fn ip() -> Result<IpAddr> {
    let connect_addr = SocketAddr::from(OutboundIpAddrLookupSocketAddr::configured_value());
    ip_impl(connect_addr).map_err(Error::NoOutboundIpAddr)
}

fn ip_impl(connect_addr: impl ToSocketAddrs) -> io::Result<IpAddr> {
    let socket = UdpSocket::bind(UNSPECIFIED_SOCKET_ADDR)?;
    socket.connect(connect_addr)?;
    let addr = socket.local_addr()?;
    Ok(addr.ip())
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn ip_lookup() {
        assert!(ip_impl("an invalid socket addr").is_err());
        assert!(ip().is_ok());

        // We would prefer to use the `locked_env_var` macro here but currently, the `habitat_core`
        // crate does not depend on `habitat_common`. If this becomes an issue, some refactoring
        // will be needed.
        std::env::set_var(OutboundIpAddrLookupSocketAddr::ENVVAR, "127.0.0.1:80");
        // localhost should route back to localhost
        assert_eq!(ip().unwrap(), Ipv4Addr::LOCALHOST);
    }
}
<|endoftext|>