<commit_before>// Test that we don't duplicate storage for futures moved around in .await, and
// for futures moved into other futures.
//
// The exact sizes can change by a few bytes (we'd like to know when they do).
// What we don't want to see is the wrong multiple of 1024 (the size of BigFut)
// being reflected in the size.
//
// See issue #59123 for a full explanation.

// ignore-wasm32-bare (sizes don't match)
// run-pass

// edition:2018

#![feature(async_await)]

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

const BIG_FUT_SIZE: usize = 1024;
struct BigFut([u8; BIG_FUT_SIZE]);

impl BigFut {
    fn new() -> Self {
        BigFut([0; BIG_FUT_SIZE])
    } }

impl Drop for BigFut {
    fn drop(&mut self) {}
}

impl Future for BigFut {
    type Output = ();

    fn poll(self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(())
    }
}

#[allow(dead_code)]
struct Joiner {
    a: Option<BigFut>,
    b: Option<BigFut>,
    c: Option<BigFut>,
}

impl Future for Joiner {
    type Output = ();

    fn poll(self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(())
    }
}

fn noop() {}

async fn single() {
    let x = BigFut::new();
    x.await;
}

async fn single_with_noop() {
    let x = BigFut::new();
    noop();
    x.await;
}

async fn joined() {
    let a = BigFut::new();
    let b = BigFut::new();
    let c = BigFut::new();

    let joiner = Joiner {
        a: Some(a),
        b: Some(b),
        c: Some(c),
    };
    joiner.await
}

async fn joined_with_noop() {
    let a = BigFut::new();
    let b = BigFut::new();
    let c = BigFut::new();

    let joiner = Joiner {
        a: Some(a),
        b: Some(b),
        c: Some(c),
    };
    noop();
    joiner.await
}

fn main() {
    assert_eq!(1028, std::mem::size_of_val(&single()));
    assert_eq!(1032, std::mem::size_of_val(&single_with_noop()));
    assert_eq!(3084, std::mem::size_of_val(&joined()));
    assert_eq!(3084, std::mem::size_of_val(&joined_with_noop()));
}
<commit_msg>test .await while holding variables of different sizes<commit_after>// Test that we don't duplicate storage for futures moved around in .await, and
// for futures moved into other futures.
//
// The exact sizes can change by a few bytes (we'd like to know when they do).
// What we don't want to see is the wrong multiple of 1024 (the size of BigFut)
// being reflected in the size.
//
// See issue #59123 for a full explanation.

// ignore-wasm32-bare (sizes don't match)
// run-pass

// edition:2018

#![feature(async_await)]

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

const BIG_FUT_SIZE: usize = 1024;
struct BigFut([u8; BIG_FUT_SIZE]);

impl BigFut {
    fn new() -> Self {
        BigFut([0; BIG_FUT_SIZE])
    } }

impl Drop for BigFut {
    fn drop(&mut self) {}
}

impl Future for BigFut {
    type Output = ();

    fn poll(self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(())
    }
}

#[allow(dead_code)]
struct Joiner {
    a: Option<BigFut>,
    b: Option<BigFut>,
    c: Option<BigFut>,
}

impl Future for Joiner {
    type Output = ();

    fn poll(self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(())
    }
}

fn noop() {}

async fn single() {
    let x = BigFut::new();
    x.await;
}

async fn single_with_noop() {
    let x = BigFut::new();
    noop();
    x.await;
}

async fn joined() {
    let a = BigFut::new();
    let b = BigFut::new();
    let c = BigFut::new();

    let joiner = Joiner {
        a: Some(a),
        b: Some(b),
        c: Some(c),
    };
    joiner.await
}

async fn joined_with_noop() {
    let a = BigFut::new();
    let b = BigFut::new();
    let c = BigFut::new();

    let joiner = Joiner {
        a: Some(a),
        b: Some(b),
        c: Some(c),
    };
    noop();
    joiner.await
}

async fn mixed_sizes() {
    let a = BigFut::new();
    let b = BigFut::new();
    let c = BigFut::new();
    let d = BigFut::new();
    let e = BigFut::new();
    let joiner = Joiner {
        a: Some(a),
        b: Some(b),
        c: Some(c),
    };

    d.await;
    e.await;
    joiner.await;
}

fn main() {
    assert_eq!(1028, std::mem::size_of_val(&single()));
    assert_eq!(1032, std::mem::size_of_val(&single_with_noop()));
    assert_eq!(3084, std::mem::size_of_val(&joined()));
    assert_eq!(3084, std::mem::size_of_val(&joined_with_noop()));
    assert_eq!(7188, std::mem::size_of_val(&mixed_sizes()));
}
<|endoftext|>


<commit_before>//! Slog atomic switching drain
//!
//! `AtomicSwitch` allows swapping drain that it wraps atomically, race-free, in
//! runtime. This can be useful eg. for turning on debug logging
//! in production.
//!
//! See [`slog` `signal.rs`
//! example](https://github.com/dpc/slog-rs/blob/master/examples/signal.rs)
#![warn(missing_docs)]

#[macro_use]
extern crate slog;
extern crate crossbeam;

use slog::*;
use std::sync::Arc;
use crossbeam::sync::ArcCell;

/// Handle to `AtomicSwitch` that controls it.
pub struct AtomicSwitchCtrl<E>(Arc<ArcCell<Box<Drain<Error=E>>>>);

/// Drain wrapping another drain, allowing atomic substitution in runtime
pub struct AtomicSwitch<E>(Arc<ArcCell<Box<Drain<Error=E>>>>);

impl<E> AtomicSwitch<E> {
    /// Wrap `drain` in `AtomicSwitch` to allow swapping it later
    ///
    /// Use `AtomicSwitch::ctrl()` to get a handle to it
    pub fn new<D: Drain<Error=E> + 'static>(drain: D) -> Self {
        AtomicSwitch::new_from_arc(Arc::new(ArcCell::new(Arc::new(Box::new(drain) as Box<Drain<Error=E>>))))
    }

    /// Create new `AtomicSwitch` from an existing `Arc<...>`
    ///
    /// See `AtomicSwitch::new()`
    pub fn new_from_arc(d: Arc<ArcCell<Box<Drain<Error=E>>>>) -> Self {
        AtomicSwitch(d)
    }

    /// Get a `AtomicSwitchCtrl` handle to control this `AtomicSwitch` drain
    pub fn ctrl(&self) -> AtomicSwitchCtrl<E> {
        AtomicSwitchCtrl(self.0.clone())
    }
}

impl<E> AtomicSwitchCtrl<E> {
    /// Get Arc to the currently wrapped drain
    pub fn get(&self) -> Arc<Box<Drain<Error=E>>> {
        self.0.get()
    }

    /// Set the current wrapped drain
    pub fn set<D: Drain<Error=E>>(&self, drain: D) {
        let _ = self.0.set(Arc::new(Box::new(drain)));
    }

    /// Swap the existing drain with a new one
    pub fn swap(&self, drain: Arc<Box<Drain<Error=E>>>) -> Arc<Box<Drain<Error=E>>> {
        self.0.set(drain)
    }

    /// Get a `AtomicSwitch` drain controlled by this `AtomicSwitchCtrl`
    pub fn drain(&self) -> AtomicSwitch<E> {
        AtomicSwitch(self.0.clone())
    }
}

impl<E> Drain for AtomicSwitch<E> {
    type Error = E;
    fn log(&self, info: &Record, logger_values: &OwnedKeyValueList) -> std::result::Result<(), E> {
        self.0.get().log(info, logger_values)
    }
}
<commit_msg>Prepare for `Send+Sync` removal from `Drain`<commit_after>//! Slog atomic switching drain
//!
//! `AtomicSwitch` allows swapping drain that it wraps atomically, race-free, in
//! runtime. This can be useful eg. for turning on debug logging
//! in production.
//!
//! See [`slog` `signal.rs`
//! example](https://github.com/dpc/slog-rs/blob/master/examples/signal.rs)
#![warn(missing_docs)]

#[macro_use]
extern crate slog;
extern crate crossbeam;

use slog::*;
use std::sync::Arc;
use crossbeam::sync::ArcCell;

/// Handle to `AtomicSwitch` that controls it.
pub struct AtomicSwitchCtrl<E>(Arc<ArcCell<Box<Drain<Error=E>+Send+Sync>>>);

/// Drain wrapping another drain, allowing atomic substitution in runtime
pub struct AtomicSwitch<E>(Arc<ArcCell<Box<Drain<Error=E>+Send+Sync>>>);

impl<E> AtomicSwitch<E> {
    /// Wrap `drain` in `AtomicSwitch` to allow swapping it later
    ///
    /// Use `AtomicSwitch::ctrl()` to get a handle to it
    pub fn new<D: Drain<Error=E> + 'static + Send+Sync>(drain: D) -> Self {
        AtomicSwitch::new_from_arc(Arc::new(ArcCell::new(Arc::new(Box::new(drain) as Box<Drain<Error=E>+Send+Sync>))))
    }

    /// Create new `AtomicSwitch` from an existing `Arc<...>`
    ///
    /// See `AtomicSwitch::new()`
    pub fn new_from_arc(d: Arc<ArcCell<Box<Drain<Error=E>+Send+Sync>>>) -> Self {
        AtomicSwitch(d)
    }

    /// Get a `AtomicSwitchCtrl` handle to control this `AtomicSwitch` drain
    pub fn ctrl(&self) -> AtomicSwitchCtrl<E> {
        AtomicSwitchCtrl(self.0.clone())
    }
}

impl<E> AtomicSwitchCtrl<E> {
    /// Get Arc to the currently wrapped drain
    pub fn get(&self) -> Arc<Box<Drain<Error=E>+Send+Sync>> {
        self.0.get()
    }

    /// Set the current wrapped drain
    pub fn set<D: Drain<Error=E>+Send+Sync>(&self, drain: D) {
        let _ = self.0.set(Arc::new(Box::new(drain)));
    }

    /// Swap the existing drain with a new one
    pub fn swap(&self, drain: Arc<Box<Drain<Error=E>+Send+Sync>>) -> Arc<Box<Drain<Error=E>+Send+Sync>> {
        self.0.set(drain)
    }

    /// Get a `AtomicSwitch` drain controlled by this `AtomicSwitchCtrl`
    pub fn drain(&self) -> AtomicSwitch<E> {
        AtomicSwitch(self.0.clone())
    }
}

impl<E> Drain for AtomicSwitch<E> {
    type Error = E;
    fn log(&self, info: &Record, logger_values: &OwnedKeyValueList) -> std::result::Result<(), E> {
        self.0.get().log(info, logger_values)
    }
}
<|endoftext|>


<commit_before>use futures::{Async, AsyncSink, Future, Poll};
use futures::stream::Stream;
use futures::sink::Sink;
use state_machine_future::RentToOwn;
use tokio_core::reactor::Handle;
use tokio_core::net::TcpStream;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_io::io::read;
use tokio_retry::{self, Retry};
use tokio_retry::strategy::{FibonacciBackoff, jitter};
use std::net::SocketAddr;
use std::io;
use std::time::Duration;
use bytes::{Bytes, BytesMut, Buf, IntoBuf};

pub type ConnectionError = io::Error;

#[derive(StateMachineFuture)]
pub enum Connection<Item: Into<Vec<u8>>, S: Stream<Item=Item,Error=io::Error>, K: Sink<SinkItem=Vec<u8>,SinkError=io::Error>> {
    #[state_machine_future(start, transitions(Connecting, Finished, Error))]
    NotConnected {
        addr: SocketAddr,
        handle: Handle,
        stream: S,
        sink: K,
        error_count: u32,
    },

    #[state_machine_future(transitions(NotConnected, Connected, Finished, Error))]
    Connecting {
        addr: SocketAddr,
        handle: Handle,
        stream: S,
        sink: K,
        tcp_future: Box<Future<Item=TcpStream,Error=io::Error>>,
    },

    #[state_machine_future(transitions(Connected, NotConnected, Finished, Error))]
    Connected {
        addr: SocketAddr,
        handle: Handle,
        stream: S,
        sink: K,
        tcp: TcpStream,
        outbound: Option<Bytes>,
        inbound: Option<Bytes>,
        inbound_inflight: u32,
    },

    #[state_machine_future(ready)]
    Finished(()),

    #[state_machine_future(error)]
    Error(ConnectionError),
}

impl <Item, S, K> Connection<Item, S, K> where Item: Into<Vec<u8>>, S: Stream<Item=Item,Error=io::Error>, K: Sink<SinkItem=Vec<u8>,SinkError=io::Error> {
    pub fn new(addr: SocketAddr, handle: Handle, stream: S, sink: K) -> ConnectionFuture<Item, S, K> {
        debug!("constructing new connection future");
        Connection::start(addr, handle, stream, sink, 0)
    }
}

impl <Item, S, K> PollConnection<Item, S, K> for Connection<Item, S, K> where Item: Into<Vec<u8>>, S: Stream<Item=Item,Error=io::Error>, K: Sink<SinkItem=Vec<u8>,SinkError=io::Error> {
    fn poll_not_connected<'a>(not_connected: &'a mut RentToOwn<'a, NotConnected<Item, S, K>>) -> Poll<AfterNotConnected<Item, S, K>, ConnectionError> {
        trace!("not connected to {}", not_connected.addr);

        let retry_strategy = FibonacciBackoff ::from_millis(1)
            .max_delay(Duration::from_millis(2000))
            .map(jitter);

        let addr = not_connected.addr.clone();
        let handle = not_connected.handle.clone();

        let tcp_future = Box::new(Retry::spawn(not_connected.handle.clone(), retry_strategy, move || {
            TcpStream::connect(&addr, &handle)
        }).map_err(|e| {
            match e {
                tokio_retry::Error::OperationError(e) => e,
                tokio_retry::Error::TimerError(e) => e,
            }
        })) as Box<Future<Item=TcpStream,Error=io::Error>>;

        let not_connected = not_connected.take();
        Ok(Async::Ready(Connecting {
            addr: not_connected.addr,
            handle: not_connected.handle,
            stream: not_connected.stream,
            sink: not_connected.sink,
            tcp_future: tcp_future,
        }.into()))
    }

    fn poll_connecting<'a>(connecting: &'a mut RentToOwn<'a, Connecting<Item, S, K>>) -> Poll<AfterConnecting<Item, S, K>, ConnectionError> {
        trace!("connecting to {}", connecting.addr);

        let tcp = try_ready!(connecting.tcp_future.poll());

        let connecting = connecting.take();
        Ok(Async::Ready(Connected {
            addr: connecting.addr,
            handle: connecting.handle,
            stream: connecting.stream,
            sink: connecting.sink,
            tcp: tcp,
            inbound: None,
            inbound_inflight: 0,
            outbound: None,
        }.into()))
    }

    fn poll_connected<'a>(connected: &'a mut RentToOwn<'a, Connected<Item, S, K>>) -> Poll<AfterConnected<Item, S, K>, ConnectionError> {
        trace!("connected to {}", connected.addr);
        // use std::io::Read;

        // let buf = try_ready!(
        //         connected
        //             .stream
        //             .by_ref()
        //             .take(1)
        //             .map(|msg| msg.into())
        //             .concat2()
        //             .select(
        //                 read(connected.tcp.by_ref(), vec![])
        //                 .map(|(tcp, buf, len)| {
        //                     buf
        //                 })
        //             )
        //             .poll()
        //         );

        let mut progress = false;

        let mut received = connected.inbound.take();
        if received.is_none() && connected.inbound_inflight == 0 {
            let mut receiving_buf = BytesMut::new();
            received = match connected.tcp.read_buf(&mut receiving_buf) {
                Ok(Async::Ready(len)) => {
                    trace!("tcpstream received {} bytes", len);
                    progress = true;
                    Some(receiving_buf.freeze())
                },
                Ok(Async::NotReady) => {
                    trace!("tcpstream not ready to read");
                    None
                },
                Err(e) => {
                    trace!("tcpstream read error: {:?}", e);
                    let connected = connected.take();

                    return Ok(Async::Ready(NotConnected {
                        addr: connected.addr,
                        handle: connected.handle,
                        stream: connected.stream,
                        sink: connected.sink,
                        error_count: 0,
                    }.into()));
                }
            };
        }

        let mut sending_buf = connected.outbound.take();
        if sending_buf.is_none() {
            sending_buf = match try!(connected.stream.poll()) {
                Async::Ready(Some(msg)) => {
                    progress = true;
                    let mut buf = Bytes::new();
                    let msg_vec: Vec<u8> = msg.into();
                    trace!("stream received new message of length {}", msg_vec.len());
                    buf.extend_from_slice(&msg_vec[..]);

                    Some(buf)
                },
                Async::Ready(None) => {
                    trace!("stream finished");
                    return Ok(Async::Ready(Finished(()).into()));
                },
                Async::NotReady => {
                    trace!("stream not ready to read");
                    None
                }
            };
        }

        if let Some(buf) = sending_buf {
            let total_len = buf.len();
            let mut buf = buf.into_buf();
            match connected.tcp.write_buf(&mut buf) {
                Ok(Async::Ready(len)) => {
                    trace!("tcpstream wrote {} of {} bytes", len, total_len);
                    progress = true;
                    if buf.has_remaining() {
                        let mut buf_remaining = Bytes::new();
                        buf_remaining.extend_from_slice(buf.bytes());
                        connected.outbound = Some(buf_remaining);
                    }
                },
                Ok(Async::NotReady) => {
                    trace!("tcpstream not ready to write");
                    let mut buf_remaining = Bytes::new();
                    buf_remaining.extend_from_slice(buf.bytes());
                    connected.outbound = Some(buf_remaining);
                },
                Err(err) => {
                    trace!("tcpstream write error: {:?}", err);
                    unimplemented!();
                }
            }
        }

        if let Some(buf) = received {
            match try!(connected.sink.start_send(buf.to_vec())) {
                AsyncSink::Ready => {
                    trace!("sink started send");
                    progress = true;
                    connected.inbound_inflight += 1;
                },
                AsyncSink::NotReady(buf) => {
                    trace!("sink not ready");
                    let mut buf_remaining = Bytes::new();
                    buf_remaining.extend_from_slice(&buf[..]);
                    connected.inbound = Some(buf_remaining);
                },
            }
        }

        if connected.inbound_inflight != 0 {
            match try!(connected.sink.poll_complete()) {
                Async::Ready(()) => {
                    trace!("sink polled complete");
                    progress = true;
                    connected.inbound_inflight = 0;
                },
                Async::NotReady => {
                    trace!("sink polled not ready");
                    //
                }
            }
        }

        if progress {
            trace!("made progress");

            let connected = connected.take();

            return Ok(Async::Ready(Connected {
                addr: connected.addr,
                handle: connected.handle,
                stream: connected.stream,
                sink: connected.sink,
                tcp: connected.tcp,
                inbound: connected.inbound,
                inbound_inflight: connected.inbound_inflight,
                outbound: connected.outbound,
            }.into()));
        } else {
            trace!("did not make progress");

            return Ok(Async::NotReady);
        }
    }

    // fn poll_connected_sending<'a>(connected: &'a mut RentToOwn<'a, ConnectedSending<Item, S>>) -> Poll<AfterConnectedSending<Item, S>, ConnectionError> {
    //     let mut buf = connected.buf.take().unwrap().into_buf();

    //     trace!("trying to write {} bytes", buf.remaining());
    //     match connected.tcp.write_buf(&mut buf) {
    //         Ok(Async::Ready(len)) => {
    //             trace!("wrote {} bytes", len);

    //             let mut connected = connected.take();

    //             if buf.has_remaining() {
    //                 let mut buf_remaining = Bytes::new();
    //                 buf_remaining.extend_from_slice(buf.bytes());
    //                 connected.buf = Some(buf_remaining);

    //                 return Ok(Async::Ready(ConnectedSending {
    //                     addr: connected.addr,
    //                     handle: connected.handle,
    //                     stream: connected.stream,
    //                     tcp: connected.tcp,
    //                     buf: connected.buf,
    //                 }.into()));
    //             } else if (connected.inbound.is_some()) {
    //                 return Ok(Async::Ready(ConnectedReceiving {
    //                     addr: connected.addr,
    //                     handle: connected.handle,
    //                     stream: connected.stream,
    //                     tcp: connected.tcp,
    //                 }.into()));
    //             } else {
    //                 return Ok(Async::Ready(Connected {
    //                     addr: connected.addr,
    //                     handle: connected.handle,
    //                     stream: connected.stream,
    //                     tcp: connected.tcp,
    //                 }.into()));
    //             }
    //         },
    //         Ok(Async::NotReady) => {
    //             trace!("not ready");
    //             let mut buf_remaining = Bytes::new();
    //             buf_remaining.extend_from_slice(buf.bytes());
    //             connected.buf = Some(buf_remaining);
    //             return Ok(Async::NotReady);
    //         },
    //         Err(err) => {
    //             debug!("error writing to socket {:?}", err);
    //             let connected = connected.take();

    //             return Ok(Async::Ready(NotConnected {
    //                 addr: connected.addr,
    //                 handle: connected.handle,
    //                 stream: connected.stream,
    //                 error_count: 0,
    //             }.into()));
    //         }
    //     }
    // }

    // fn poll_connected_receiving<'a>(connected: &'a mut RentToOwn<'a, ConnectedReceiving<Item, S>>) -> Poll<AfterConnectedReceiving<Item, S>, ConnectionError> {

    // }
}
<commit_msg>include error counts, clean up code<commit_after>use futures::{Async, AsyncSink, Future, Poll};
use futures::stream::Stream;
use futures::sink::Sink;
use state_machine_future::RentToOwn;
use tokio_core::reactor::Handle;
use tokio_core::net::TcpStream;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_retry::{self, Retry};
use tokio_retry::strategy::{FibonacciBackoff, jitter};
use std::ops::Deref;
use std::net::SocketAddr;
use std::io;
use std::time::Duration;
use bytes::{Bytes, BytesMut, Buf, IntoBuf};

pub type ConnectionError = io::Error;

#[derive(StateMachineFuture)]
pub enum Connection<Item: Into<Vec<u8>>, S: Stream<Item=Item,Error=io::Error>, K: Sink<SinkItem=Vec<u8>,SinkError=io::Error>> {
    #[state_machine_future(start, transitions(Connecting, Finished, Error))]
    NotConnected {
        addr: SocketAddr,
        handle: Handle,
        stream: S,
        sink: K,
        error_count: u32,
    },

    #[state_machine_future(transitions(NotConnected, Connected, Finished, Error))]
    Connecting {
        addr: SocketAddr,
        handle: Handle,
        stream: S,
        sink: K,
        tcp_future: Box<Future<Item=TcpStream,Error=io::Error>>,
        error_count: u32,
    },

    #[state_machine_future(transitions(Connected, NotConnected, Finished, Error))]
    Connected {
        addr: SocketAddr,
        handle: Handle,
        stream: S,
        sink: K,
        tcp: TcpStream,
        outbound: Option<Bytes>,
        inbound: Option<Bytes>,
        inbound_inflight: u32,
        error_count: u32,
    },

    #[state_machine_future(ready)]
    Finished(()),

    #[state_machine_future(error)]
    Error(ConnectionError),
}

impl <Item, S, K> Connection<Item, S, K> where Item: Into<Vec<u8>>, S: Stream<Item=Item,Error=io::Error>, K: Sink<SinkItem=Vec<u8>,SinkError=io::Error> {
    pub fn new(addr: SocketAddr, handle: Handle, stream: S, sink: K) -> ConnectionFuture<Item, S, K> {
        debug!("constructing new connection future");
        Connection::start(addr, handle, stream, sink, 0)
    }
}

impl <Item, S, K> PollConnection<Item, S, K> for Connection<Item, S, K> where Item: Into<Vec<u8>>, S: Stream<Item=Item,Error=io::Error>, K: Sink<SinkItem=Vec<u8>,SinkError=io::Error> {
    fn poll_not_connected<'a>(not_connected: &'a mut RentToOwn<'a, NotConnected<Item, S, K>>) -> Poll<AfterNotConnected<Item, S, K>, ConnectionError> {
        trace!("not connected to {}", not_connected.addr);

        let retry_strategy = FibonacciBackoff ::from_millis(1)
            .max_delay(Duration::from_millis(2000))
            .map(jitter);

        let addr = not_connected.addr.clone();
        let handle = not_connected.handle.clone();

        let tcp_future = Box::new(Retry::spawn(not_connected.handle.clone(), retry_strategy, move || {
            TcpStream::connect(&addr, &handle)
        }).map_err(|e| {
            match e {
                tokio_retry::Error::OperationError(e) => e,
                tokio_retry::Error::TimerError(e) => e,
            }
        })) as Box<Future<Item=TcpStream,Error=io::Error>>;

        let not_connected = not_connected.take();
        Ok(Async::Ready(Connecting {
            addr: not_connected.addr,
            handle: not_connected.handle,
            stream: not_connected.stream,
            sink: not_connected.sink,
            tcp_future: tcp_future,
            error_count: not_connected.error_count,
        }.into()))
    }

    fn poll_connecting<'a>(connecting: &'a mut RentToOwn<'a, Connecting<Item, S, K>>) -> Poll<AfterConnecting<Item, S, K>, ConnectionError> {
        trace!("connecting to {}", connecting.addr);

        let tcp = try_ready!(connecting.tcp_future.poll());

        let connecting = connecting.take();
        Ok(Async::Ready(Connected {
            addr: connecting.addr,
            handle: connecting.handle,
            stream: connecting.stream,
            sink: connecting.sink,
            tcp: tcp,
            inbound: None,
            inbound_inflight: 0,
            outbound: None,
            error_count: connecting.error_count,
        }.into()))
    }

    fn poll_connected<'a>(connected: &'a mut RentToOwn<'a, Connected<Item, S, K>>) -> Poll<AfterConnected<Item, S, K>, ConnectionError> {
        trace!("connected to {}", connected.addr);

        // we make progress if any poll returns Ready
        // if we make progress then we return a state or state change
        // otherwise return NotReady
        let mut progress = false;

        let mut received = connected.inbound.take();
        if received.is_none() && connected.inbound_inflight == 0 {
            let mut receiving_buf = BytesMut::new();
            received = match connected.tcp.read_buf(&mut receiving_buf) {
                Ok(Async::Ready(len)) => {
                    trace!("tcpstream received {} bytes", len);
                    progress = true;
                    Some(receiving_buf.freeze())
                },
                Ok(Async::NotReady) => {
                    trace!("tcpstream not ready to read");
                    None
                },
                Err(e) => {
                    trace!("tcpstream read error: {:?}", e);
                    let connected = connected.take();

                    return Ok(Async::Ready(NotConnected {
                        addr: connected.addr,
                        handle: connected.handle,
                        stream: connected.stream,
                        sink: connected.sink,
                        error_count: connected.error_count + 1,
                    }.into()));
                }
            };
        }

        let mut sending_buf = connected.outbound.take();
        if sending_buf.is_none() {
            sending_buf = match try!(connected.stream.poll()) {
                Async::Ready(Some(msg)) => {
                    progress = true;
                    let mut buf = Some(Bytes::new());
                    let msg_vec: Vec<u8> = msg.into();
                    trace!("stream received new message of length {}", msg_vec.len());
                    set_bytes(&mut buf, msg_vec);

                    buf
                },
                Async::Ready(None) => {
                    trace!("stream finished");
                    return Ok(Async::Ready(Finished(()).into()));
                },
                Async::NotReady => {
                    trace!("stream not ready to read");
                    None
                }
            };
        }

        if let Some(buf) = sending_buf {
            let total_len = buf.len();
            let mut buf = buf.into_buf();
            match connected.tcp.write_buf(&mut buf) {
                Ok(Async::Ready(len)) => {
                    trace!("tcpstream wrote {} of {} bytes", len, total_len);
                    progress = true;
                    if buf.has_remaining() {
                        set_bytes(&mut connected.outbound, buf.bytes());
                    }
                },
                Ok(Async::NotReady) => {
                    trace!("tcpstream not ready to write");
                    set_bytes(&mut connected.outbound, buf.bytes());
                },
                Err(err) => {
                    trace!("tcpstream write error: {:?}", err);
                    set_bytes(&mut connected.outbound, buf.bytes());

                    let connected = connected.take();

                    return Ok(Async::Ready(NotConnected {
                        addr: connected.addr,
                        handle: connected.handle,
                        stream: connected.stream,
                        sink: connected.sink,
                        error_count: connected.error_count + 1,
                    }.into()));
                }
            }
        }

        if let Some(buf) = received {
            match try!(connected.sink.start_send(buf.to_vec())) {
                AsyncSink::Ready => {
                    trace!("sink started send");
                    progress = true;
                    connected.inbound_inflight += 1;
                },
                AsyncSink::NotReady(buf) => {
                    trace!("sink not ready");
                    set_bytes(&mut connected.inbound, buf);
                },
            }
        }

        if connected.inbound_inflight != 0 {
            match try!(connected.sink.poll_complete()) {
                Async::Ready(()) => {
                    trace!("sink polled complete");
                    progress = true;
                    connected.inbound_inflight = 0;
                },
                Async::NotReady => {
                    trace!("sink polled not ready");
                    // do nothing
                }
            }
        }

        if progress {
            trace!("made progress");

            let connected = connected.take();

            return Ok(Async::Ready(Connected {
                addr: connected.addr,
                handle: connected.handle,
                stream: connected.stream,
                sink: connected.sink,
                tcp: connected.tcp,
                inbound: connected.inbound,
                inbound_inflight: connected.inbound_inflight,
                outbound: connected.outbound,
                error_count: connected.error_count,
            }.into()));
        } else {
            trace!("did not make progress");

            return Ok(Async::NotReady);
        }
    }
}

fn set_bytes<'a, B>(maybe_buffer: &mut Option<Bytes>, content: B) where B: Deref<Target=[u8]> {
     let mut buf_remaining = Bytes::new();
    buf_remaining.extend_from_slice(&content);
    *maybe_buffer = Some(buf_remaining);
}<|endoftext|>


<commit_before>use std::os::unix::io::RawFd;
use std::collections::HashMap;
use std::slice;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::os::unix::io::AsRawFd;
use nix::sys::event::{ev_set, kqueue, kevent, KEvent, EventFilter, EventFlag, FilterFlag};
use libc::{uintptr_t, intptr_t};
use std::io::Result;
use event::Event;
use notification::Notification;
use user_event::UserEvent;
use nix_err_to_io_err;

type UserData = intptr_t;

static KQUEUE_EVENT_SIZE: usize = 1024;

pub struct KernelPoller {
    kqueue: RawFd,
    registrar: KernelRegistrar,
    eventlist: Vec<KEvent>,
    notifications: HashMap<RawFd, Notification>
}

impl KernelPoller {
    pub fn new() -> Result<KernelPoller> {
        let kq = kqueue().map_err(nix_err_to_io_err)?;
        let registrations = Arc::new(AtomicUsize::new(1));
        Ok(KernelPoller {
            kqueue: kq,
            registrar: KernelRegistrar::new(kq, registrations),
            eventlist: Vec::with_capacity(KQUEUE_EVENT_SIZE),
            notifications: HashMap::with_capacity(KQUEUE_EVENT_SIZE)
        })
    }

    // This will always succeed. We implement it this way for api compatibility with epoll.
    pub fn get_registrar(&self) -> Result<KernelRegistrar> {
        Ok(self.registrar.clone())
    }

    // Wait for kevents and return a list of Notifications. Coalesce reads and writes for the same
    // socket into a single notification. If only a read or a write event for a given socket is
    // present in the eventlist, check the registration to see if there is another kevent registered
    //  and remove it if so. We do this removal to prevent aliasing a pointer to the same
    // registration structure.
    pub fn wait(&mut self, timeout_ms: usize) -> Result<Vec<Notification>> {

        // Create a buffer to read events into
        let dst = unsafe {
            slice::from_raw_parts_mut(self.eventlist.as_mut_ptr(), self.eventlist.capacity())
        };

        let count = kevent(self.kqueue, &[], dst, timeout_ms).map_err(nix_err_to_io_err)?;

        // Set the length of the vector to the number of events that was returned by kevent
        unsafe { self.eventlist.set_len(count); }

        self.coalesce_events();
        Ok(self.notifications.drain().map(|(_, v)| v).collect())
    }

    // Combine read and write events for the same socket into a single notification.
    fn coalesce_events(&mut self) {
        for e in self.eventlist.drain(..) {
            let event = event_from_filter(e.filter());
            let new_notification = Notification {
                id: e.udata() as usize,
                event: event.clone()
            };

            let mut notification = self.notifications.entry(e.ident() as RawFd)
                                                     .or_insert(new_notification);
            if notification.event != event {
                notification.event = Event::Both
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct KernelRegistrar {
    kqueue: RawFd,
    total_registrations: Arc<AtomicUsize>
}

impl KernelRegistrar {
    // Explicitly not public. KernelRegistrar's are tied to KernelPollers and are retreived via
    // calls to poller.get_registrar().
    fn new(kq: RawFd, registrations: Arc<AtomicUsize>) -> KernelRegistrar {
        KernelRegistrar {
            kqueue: kq,
            total_registrations: registrations
        }
    }

    // This will always succeed. We implement this to provide api compatibility with epoll.
    pub fn try_clone(&self) -> Result<KernelRegistrar> {
        Ok(self.clone())
    }

    pub fn register<T: AsRawFd>(&self, sock: &T, event: Event) -> Result<usize> {
        let sock_fd = sock.as_raw_fd();
        let id = self.total_registrations.fetch_add(1, Ordering::SeqCst);
        let changes = make_changelist(sock_fd, event, id as UserData);
        kevent(self.kqueue, &changes, &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(id)
    }

    pub fn reregister<T: AsRawFd>(&self, id: usize, sock: &T, event: Event) -> Result<()> {
        let sock_fd = sock.as_raw_fd();
        let changes = make_changelist(sock_fd, event, id as UserData);
        kevent(self.kqueue, &changes, &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    pub fn deregister<T: AsRawFd>(&self, sock: &T) -> Result<()> {
        let sock_fd = sock.as_raw_fd();
        let mut changes = make_changelist(sock_fd, Event::Both, 0);
        for e in changes.iter_mut() {
            set_flags(e, EventFlag::EV_DELETE);
        }
        // Just ignore errors because, one of the events may not be present, but the deregister
        // signature ignores that fact. At this point, ownership of the socket is taken so it's
        // irrelevant anyway.
        let _ = kevent(self.kqueue, &changes, &mut[], 0);
        Ok(())
    }

    pub fn register_user_event(&mut self) -> Result<UserEvent> {
        let id = self.total_registrations.fetch_add(1, Ordering::SeqCst);
        let changes = vec![make_user_event(id)];
        kevent(self.kqueue, &changes, &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(UserEvent {id: id, registrar: self.clone()})
    }

    pub fn trigger_user_event(&self, event: &UserEvent) -> Result<()> {
        let mut e = make_user_event(event.get_id());
        set_flags(&mut e, EventFlag::EV_ENABLE);
        set_fflags(&mut e, FilterFlag::NOTE_TRIGGER);
        kevent(self.kqueue, &vec![e], &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    pub fn clear_user_event(&self, event: &UserEvent) -> Result<()> {
        let mut user_event = make_user_event(event.get_id());
        set_flags(&mut user_event, EventFlag::EV_DISABLE);
        kevent(self.kqueue, &vec![user_event], &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    pub fn deregister_user_event(&self, event_id: usize) -> Result<()> {
        let mut user_event = make_user_event(event_id);
        set_flags(&mut user_event, EventFlag::EV_DELETE);
        kevent(self.kqueue, &vec![user_event], &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    pub fn set_timeout(&self, timeout: usize) -> Result<usize> {
        self.set_timer(timeout, false)
    }

    pub fn set_interval(&self, timeout: usize) -> Result<usize> {
        self.set_timer(timeout, true)
    }

    pub fn cancel_timeout(&self, timer_id: usize) -> Result<()> {
        let mut e = make_timer(timer_id, 0, false);
        set_flags(&mut e, EventFlag::EV_DELETE);
        kevent(self.kqueue, &vec![e], &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    fn set_timer(&self, timeout: usize, recurring: bool) -> Result<usize> {
        let id = self.total_registrations.fetch_add(1, Ordering::SeqCst);
        let changes = vec![make_timer(id, timeout, recurring)];
        kevent(self.kqueue, &changes, &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(id)
    }
}

fn event_from_filter(filter: EventFilter) -> Event {
    // TODO: Change Event to allow returning TIMER events instead of marking them READ
    if filter == EventFilter::EVFILT_READ ||
       filter == EventFilter::EVFILT_TIMER ||
       filter == EventFilter::EVFILT_USER {
        Event::Read
    } else {
        Event::Write
    }
}

// Each event in kqueue must have its own filter. In other words, there are seperate events for
// reads and writes on the same socket. We create the proper number of KEvents based on the enum
// variant in the `event` paramter.
fn make_changelist(sock_fd: RawFd, event: Event, user_data: UserData) -> Vec<KEvent> {
    let mut ev = KEvent::new(
        sock_fd as uintptr_t,
        EventFilter::EVFILT_READ,
        EventFlag::EV_ADD | EventFlag::EV_CLEAR,
        FilterFlag::empty(),
        0,
        user_data
    );

    match event {
        Event::Read => vec![ev],
        Event::Write => {
            set_filter(&mut ev, EventFilter::EVFILT_WRITE);
            vec![KEvent::new(ev.ident(), EventFilter::EVFILT_WRITE, ev.flags(), ev.fflags(), ev.data(), ev.udata())]
        },
        Event::Both => vec![ev, KEvent::new(ev.ident(), EventFilter::EVFILT_WRITE, ev.flags(), ev.fflags(), ev.data(), ev.udata())]
    }
}

fn make_user_event(id: usize) -> KEvent {
    KEvent::new(
        id as uintptr_t,
        EventFilter::EVFILT_USER,
        EventFlag::EV_ADD | EventFlag::EV_CLEAR | EventFlag::EV_ENABLE,
        FilterFlag::empty(),
        0,
        id as UserData
    )
}

fn set_filter(e: &mut KEvent, filter: EventFilter) {
    let ident = e.ident();
    let flags = e.flags();
    let fflags = e.fflags();
    let udata = e.udata();
    ev_set(e, ident, filter, flags, fflags, udata);
}

fn set_flags(e: &mut KEvent, flags: EventFlag) {
    let ident = e.ident();
    let filter = e.filter();
    let fflags = e.fflags();
    let udata = e.udata();
    ev_set(e, ident, filter, flags, fflags, udata);
}

fn set_fflags(e: &mut KEvent, fflags: FilterFlag) {
    let ident = e.ident();
    let filter = e.filter();
    let flags = e.flags();
    let udata = e.udata();
    ev_set(e, ident, filter, flags, fflags, udata);
}

fn make_timer(id: usize, timeout: usize, recurring: bool) -> KEvent {
    let mut flags = EventFlag::EV_ADD;
    if !recurring {
        flags |= EventFlag::EV_ONESHOT;
    }
    let ev = KEvent::new(
        id as uintptr_t,
        EventFilter::EVFILT_TIMER,
        flags,
        FilterFlag::empty(), // timeouts are in ms by default
        timeout as intptr_t,
        id as UserData
    );
    ev
}
<commit_msg>Remove unnecessarily constructed kqueue write event<commit_after>use std::os::unix::io::RawFd;
use std::collections::HashMap;
use std::slice;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::os::unix::io::AsRawFd;
use nix::sys::event::{ev_set, kqueue, kevent, KEvent, EventFilter, EventFlag, FilterFlag};
use libc::{uintptr_t, intptr_t};
use std::io::Result;
use event::Event;
use notification::Notification;
use user_event::UserEvent;
use nix_err_to_io_err;

type UserData = intptr_t;

static KQUEUE_EVENT_SIZE: usize = 1024;

pub struct KernelPoller {
    kqueue: RawFd,
    registrar: KernelRegistrar,
    eventlist: Vec<KEvent>,
    notifications: HashMap<RawFd, Notification>
}

impl KernelPoller {
    pub fn new() -> Result<KernelPoller> {
        let kq = kqueue().map_err(nix_err_to_io_err)?;
        let registrations = Arc::new(AtomicUsize::new(1));
        Ok(KernelPoller {
            kqueue: kq,
            registrar: KernelRegistrar::new(kq, registrations),
            eventlist: Vec::with_capacity(KQUEUE_EVENT_SIZE),
            notifications: HashMap::with_capacity(KQUEUE_EVENT_SIZE)
        })
    }

    // This will always succeed. We implement it this way for api compatibility with epoll.
    pub fn get_registrar(&self) -> Result<KernelRegistrar> {
        Ok(self.registrar.clone())
    }

    // Wait for kevents and return a list of Notifications. Coalesce reads and writes for the same
    // socket into a single notification. If only a read or a write event for a given socket is
    // present in the eventlist, check the registration to see if there is another kevent registered
    //  and remove it if so. We do this removal to prevent aliasing a pointer to the same
    // registration structure.
    pub fn wait(&mut self, timeout_ms: usize) -> Result<Vec<Notification>> {

        // Create a buffer to read events into
        let dst = unsafe {
            slice::from_raw_parts_mut(self.eventlist.as_mut_ptr(), self.eventlist.capacity())
        };

        let count = kevent(self.kqueue, &[], dst, timeout_ms).map_err(nix_err_to_io_err)?;

        // Set the length of the vector to the number of events that was returned by kevent
        unsafe { self.eventlist.set_len(count); }

        self.coalesce_events();
        Ok(self.notifications.drain().map(|(_, v)| v).collect())
    }

    // Combine read and write events for the same socket into a single notification.
    fn coalesce_events(&mut self) {
        for e in self.eventlist.drain(..) {
            let event = event_from_filter(e.filter());
            let new_notification = Notification {
                id: e.udata() as usize,
                event: event.clone()
            };

            let mut notification = self.notifications.entry(e.ident() as RawFd)
                                                     .or_insert(new_notification);
            if notification.event != event {
                notification.event = Event::Both
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct KernelRegistrar {
    kqueue: RawFd,
    total_registrations: Arc<AtomicUsize>
}

impl KernelRegistrar {
    // Explicitly not public. KernelRegistrar's are tied to KernelPollers and are retreived via
    // calls to poller.get_registrar().
    fn new(kq: RawFd, registrations: Arc<AtomicUsize>) -> KernelRegistrar {
        KernelRegistrar {
            kqueue: kq,
            total_registrations: registrations
        }
    }

    // This will always succeed. We implement this to provide api compatibility with epoll.
    pub fn try_clone(&self) -> Result<KernelRegistrar> {
        Ok(self.clone())
    }

    pub fn register<T: AsRawFd>(&self, sock: &T, event: Event) -> Result<usize> {
        let sock_fd = sock.as_raw_fd();
        let id = self.total_registrations.fetch_add(1, Ordering::SeqCst);
        let changes = make_changelist(sock_fd, event, id as UserData);
        kevent(self.kqueue, &changes, &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(id)
    }

    pub fn reregister<T: AsRawFd>(&self, id: usize, sock: &T, event: Event) -> Result<()> {
        let sock_fd = sock.as_raw_fd();
        let changes = make_changelist(sock_fd, event, id as UserData);
        kevent(self.kqueue, &changes, &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    pub fn deregister<T: AsRawFd>(&self, sock: &T) -> Result<()> {
        let sock_fd = sock.as_raw_fd();
        let mut changes = make_changelist(sock_fd, Event::Both, 0);
        for e in changes.iter_mut() {
            set_flags(e, EventFlag::EV_DELETE);
        }
        // Just ignore errors because, one of the events may not be present, but the deregister
        // signature ignores that fact. At this point, ownership of the socket is taken so it's
        // irrelevant anyway.
        let _ = kevent(self.kqueue, &changes, &mut[], 0);
        Ok(())
    }

    pub fn register_user_event(&mut self) -> Result<UserEvent> {
        let id = self.total_registrations.fetch_add(1, Ordering::SeqCst);
        let changes = vec![make_user_event(id)];
        kevent(self.kqueue, &changes, &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(UserEvent {id: id, registrar: self.clone()})
    }

    pub fn trigger_user_event(&self, event: &UserEvent) -> Result<()> {
        let mut e = make_user_event(event.get_id());
        set_flags(&mut e, EventFlag::EV_ENABLE);
        set_fflags(&mut e, FilterFlag::NOTE_TRIGGER);
        kevent(self.kqueue, &vec![e], &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    pub fn clear_user_event(&self, event: &UserEvent) -> Result<()> {
        let mut user_event = make_user_event(event.get_id());
        set_flags(&mut user_event, EventFlag::EV_DISABLE);
        kevent(self.kqueue, &vec![user_event], &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    pub fn deregister_user_event(&self, event_id: usize) -> Result<()> {
        let mut user_event = make_user_event(event_id);
        set_flags(&mut user_event, EventFlag::EV_DELETE);
        kevent(self.kqueue, &vec![user_event], &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    pub fn set_timeout(&self, timeout: usize) -> Result<usize> {
        self.set_timer(timeout, false)
    }

    pub fn set_interval(&self, timeout: usize) -> Result<usize> {
        self.set_timer(timeout, true)
    }

    pub fn cancel_timeout(&self, timer_id: usize) -> Result<()> {
        let mut e = make_timer(timer_id, 0, false);
        set_flags(&mut e, EventFlag::EV_DELETE);
        kevent(self.kqueue, &vec![e], &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(())
    }

    fn set_timer(&self, timeout: usize, recurring: bool) -> Result<usize> {
        let id = self.total_registrations.fetch_add(1, Ordering::SeqCst);
        let changes = vec![make_timer(id, timeout, recurring)];
        kevent(self.kqueue, &changes, &mut[], 0).map_err(nix_err_to_io_err)?;
        Ok(id)
    }
}

fn event_from_filter(filter: EventFilter) -> Event {
    // TODO: Change Event to allow returning TIMER events instead of marking them READ
    if filter == EventFilter::EVFILT_READ ||
       filter == EventFilter::EVFILT_TIMER ||
       filter == EventFilter::EVFILT_USER {
        Event::Read
    } else {
        Event::Write
    }
}

// Each event in kqueue must have its own filter. In other words, there are seperate events for
// reads and writes on the same socket. We create the proper number of KEvents based on the enum
// variant in the `event` paramter.
fn make_changelist(sock_fd: RawFd, event: Event, user_data: UserData) -> Vec<KEvent> {
    let mut ev = KEvent::new(
        sock_fd as uintptr_t,
        EventFilter::EVFILT_READ,
        EventFlag::EV_ADD | EventFlag::EV_CLEAR,
        FilterFlag::empty(),
        0,
        user_data
    );

    match event {
        Event::Read => vec![ev],
        Event::Write => {
            set_filter(&mut ev, EventFilter::EVFILT_WRITE);
            vec![ev]
        },
        Event::Both => vec![ev, KEvent::new(ev.ident(), EventFilter::EVFILT_WRITE, ev.flags(), ev.fflags(), ev.data(), ev.udata())]
    }
}

fn make_user_event(id: usize) -> KEvent {
    KEvent::new(
        id as uintptr_t,
        EventFilter::EVFILT_USER,
        EventFlag::EV_ADD | EventFlag::EV_CLEAR | EventFlag::EV_ENABLE,
        FilterFlag::empty(),
        0,
        id as UserData
    )
}

fn set_filter(e: &mut KEvent, filter: EventFilter) {
    let ident = e.ident();
    let flags = e.flags();
    let fflags = e.fflags();
    let udata = e.udata();
    ev_set(e, ident, filter, flags, fflags, udata);
}

fn set_flags(e: &mut KEvent, flags: EventFlag) {
    let ident = e.ident();
    let filter = e.filter();
    let fflags = e.fflags();
    let udata = e.udata();
    ev_set(e, ident, filter, flags, fflags, udata);
}

fn set_fflags(e: &mut KEvent, fflags: FilterFlag) {
    let ident = e.ident();
    let filter = e.filter();
    let flags = e.flags();
    let udata = e.udata();
    ev_set(e, ident, filter, flags, fflags, udata);
}

fn make_timer(id: usize, timeout: usize, recurring: bool) -> KEvent {
    let mut flags = EventFlag::EV_ADD;
    if !recurring {
        flags |= EventFlag::EV_ONESHOT;
    }
    let ev = KEvent::new(
        id as uintptr_t,
        EventFilter::EVFILT_TIMER,
        flags,
        FilterFlag::empty(), // timeouts are in ms by default
        timeout as intptr_t,
        id as UserData
    );
    ev
}
<|endoftext|>


<commit_before>use num::Zero;
use strum_macros::{Display};
use std::marker::PhantomData;
use super::arch;
use super::arch::{Arch};

use failure::{Error, Fail};

#[derive(Debug, Fail)]
pub enum LoaderError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
}

#[derive(Display, Clone, Copy)]
pub enum FileFormat {
    Raw,  // shellcode
    PE,
}

#[derive(Display, Clone, Copy)]
pub enum Platform {
    Windows,
}

pub struct Section<A: Arch> {
    pub addr: A::RVA,
    pub buf: Vec<u8>,
    // TODO
    pub perms: u8,
    pub name: String,
}

impl <A: Arch> Section<A> {
    pub fn contains(self: &Section<A>, rva: A::RVA) -> bool {
        if rva < self.addr {
            return false;
        }

        if let Some(max) = arch::rva_add_usize::<A>(self.addr, self.buf.len()) {
            if rva >= max {
                return false
            }
        } else {
            return false
        }

        true
    }
}

pub struct LoadedModule<A: Arch> {
    pub base_address: A::VA,
    pub sections: Vec<Section<A>>,
}

pub trait Loader<A: Arch> {
    /// Fetch the number of bits for a pointer in this architecture.
    fn get_arch(&self) -> u8;
    fn get_plat(&self) -> Platform;
    fn get_file_format(&self) -> FileFormat;

    fn get_name(&self) -> String {
        return format!("{}/{}/{}", self.get_plat(), self.get_arch(), self.get_file_format());
    }

    /// Returns True if this Loader knows how to load the given bytes.
    fn taste(&self, buf: &[u8]) -> bool;

    /// Load the given bytes into a Module.
    fn load(&self, buf: &[u8]) -> Result<LoadedModule<A>, Error>;
}

pub struct ShellcodeLoader<A: Arch> {
    plat: Platform,
    // ShellcodeLoader must have a type parameter for it
    //  to implement Loader<A>.
    // however, it doesn't actually use this type itself.
    // so, we use a phantom data marker which has zero type,
    //  to ensure there is not an unused type parameter,
    //  which is a compile error.
    _phantom: PhantomData<A>,
}

impl<A: Arch> ShellcodeLoader<A> {
    pub fn new(plat: Platform) -> ShellcodeLoader<A> {
        ShellcodeLoader {
            plat,
            _phantom: PhantomData{},
        }
    }
}

impl <A: Arch> Loader<A> for ShellcodeLoader<A> {
    fn get_arch(&self) -> u8 {
        A::get_bits()
    }

    fn get_plat(&self) -> Platform {
        self.plat
    }

    fn get_file_format(&self) -> FileFormat {
        FileFormat::Raw
    }

    fn taste(&self, _buf: &[u8]) -> bool {
        // we can load anything as shellcode
        true
    }

    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader::*;
    ///
    /// let loader = lancelot::loader::ShellcodeLoader::<Arch32>::new(Platform::Windows);
    /// loader.load(b"MZ\x90\x00")
    ///   .map(|module| {
    ///     assert_eq!(module.base_address,     0x0);
    ///     assert_eq!(module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    fn load(&self, buf: &[u8]) -> Result<LoadedModule<A>, Error> {
        Ok(LoadedModule::<A>{
            base_address: A::VA::zero(),
            sections: vec![
                Section::<A> {
                    addr: A::RVA::zero(),
                    buf: buf.to_vec(),
                    perms: 0x0, // TODO
                    name: "raw".to_string(),
                }
            ]
        })
    }
}

pub fn default_loaders<A: Arch + 'static>() -> Vec<Box<dyn Loader<A>>> {
    // we might like these to come from a lazy_static global,
    //  however, then these have to be Sync.
    // I'm not sure if that's a good idea yet.
    let mut loaders: Vec<Box<dyn Loader<A>>> = vec![];
    // the order here matters!
    // the default `load` routine will pick the first matching loader,
    //  so the earlier entries here have higher precedence.

    loaders.push(Box::new(ShellcodeLoader::<A>::new(Platform::Windows)));

    loaders
}

/// Find the loaders that support loading the given sample.
///
/// The result is an iterator so that the caller can use the first
///  matching result without waiting for all Loaders to taste the bytes.
///
/// Loaders are tasted in the order defined in `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// match taste::<Arch32>(b"\xEB\xFE").nth(0) {
///   Some(loader) => assert_eq!(loader.get_name(), "Windows/32/Raw"),
///   None => panic!("no matching loaders"),
/// };
/// ```
pub fn taste<A: Arch + 'static>(buf: &[u8]) -> impl Iterator<Item=Box<dyn Loader<A>>> {
    default_loaders::<A>()
        .into_iter()
        .filter(move |loader| loader.taste(buf))
}

/// Load the given sample as a 32-bit module using the first matching
///  loader from `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// load::<Arch32>(b"\xEB\xFE")
///   .map(|(loader, module)| {
///     assert_eq!(loader.get_name(),       "Windows/32/Raw");
///     assert_eq!(module.base_address,     0x0);
///     assert_eq!(module.sections[0].name, "raw");
///   })
///   .map_err(|e| panic!(e));
/// ```
pub fn load<A: Arch + 'static>(buf: &[u8]) -> Result<(Box<dyn Loader<A>>, LoadedModule<A>), Error> {
    match taste::<A>(buf).nth(0) {
        Some(loader) => {
            loader.load(buf).map(|module|
                (loader, module)
            )
        },
        None => Err(LoaderError::NotSupported.into()),
    }
}
<commit_msg>loader: mock out `Section.is_executable()` accessor<commit_after>use num::Zero;
use strum_macros::{Display};
use std::marker::PhantomData;
use super::arch;
use super::arch::{Arch};

use failure::{Error, Fail};

#[derive(Debug, Fail)]
pub enum LoaderError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
}

#[derive(Display, Clone, Copy)]
pub enum FileFormat {
    Raw,  // shellcode
    PE,
}

#[derive(Display, Clone, Copy)]
pub enum Platform {
    Windows,
}

pub struct Section<A: Arch> {
    pub addr: A::RVA,
    pub buf: Vec<u8>,
    // TODO
    pub perms: u8,
    pub name: String,
}

impl <A: Arch> Section<A> {
    pub fn contains(self: &Section<A>, rva: A::RVA) -> bool {
        if rva < self.addr {
            return false;
        }

        if let Some(max) = arch::rva_add_usize::<A>(self.addr, self.buf.len()) {
            if rva >= max {
                return false
            }
        } else {
            return false
        }

        true
    }

    pub fn is_executable(&self) -> bool {
        // TODO
        true
    }
}

pub struct LoadedModule<A: Arch> {
    pub base_address: A::VA,
    pub sections: Vec<Section<A>>,
}

pub trait Loader<A: Arch> {
    /// Fetch the number of bits for a pointer in this architecture.
    fn get_arch(&self) -> u8;
    fn get_plat(&self) -> Platform;
    fn get_file_format(&self) -> FileFormat;

    fn get_name(&self) -> String {
        return format!("{}/{}/{}", self.get_plat(), self.get_arch(), self.get_file_format());
    }

    /// Returns True if this Loader knows how to load the given bytes.
    fn taste(&self, buf: &[u8]) -> bool;

    /// Load the given bytes into a Module.
    fn load(&self, buf: &[u8]) -> Result<LoadedModule<A>, Error>;
}

pub struct ShellcodeLoader<A: Arch> {
    plat: Platform,
    // ShellcodeLoader must have a type parameter for it
    //  to implement Loader<A>.
    // however, it doesn't actually use this type itself.
    // so, we use a phantom data marker which has zero type,
    //  to ensure there is not an unused type parameter,
    //  which is a compile error.
    _phantom: PhantomData<A>,
}

impl<A: Arch> ShellcodeLoader<A> {
    pub fn new(plat: Platform) -> ShellcodeLoader<A> {
        ShellcodeLoader {
            plat,
            _phantom: PhantomData{},
        }
    }
}

impl <A: Arch> Loader<A> for ShellcodeLoader<A> {
    fn get_arch(&self) -> u8 {
        A::get_bits()
    }

    fn get_plat(&self) -> Platform {
        self.plat
    }

    fn get_file_format(&self) -> FileFormat {
        FileFormat::Raw
    }

    fn taste(&self, _buf: &[u8]) -> bool {
        // we can load anything as shellcode
        true
    }

    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader::*;
    ///
    /// let loader = lancelot::loader::ShellcodeLoader::<Arch32>::new(Platform::Windows);
    /// loader.load(b"MZ\x90\x00")
    ///   .map(|module| {
    ///     assert_eq!(module.base_address,     0x0);
    ///     assert_eq!(module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    fn load(&self, buf: &[u8]) -> Result<LoadedModule<A>, Error> {
        Ok(LoadedModule::<A>{
            base_address: A::VA::zero(),
            sections: vec![
                Section::<A> {
                    addr: A::RVA::zero(),
                    buf: buf.to_vec(),
                    perms: 0x0, // TODO
                    name: "raw".to_string(),
                }
            ]
        })
    }
}

pub fn default_loaders<A: Arch + 'static>() -> Vec<Box<dyn Loader<A>>> {
    // we might like these to come from a lazy_static global,
    //  however, then these have to be Sync.
    // I'm not sure if that's a good idea yet.
    let mut loaders: Vec<Box<dyn Loader<A>>> = vec![];
    // the order here matters!
    // the default `load` routine will pick the first matching loader,
    //  so the earlier entries here have higher precedence.

    loaders.push(Box::new(ShellcodeLoader::<A>::new(Platform::Windows)));

    loaders
}

/// Find the loaders that support loading the given sample.
///
/// The result is an iterator so that the caller can use the first
///  matching result without waiting for all Loaders to taste the bytes.
///
/// Loaders are tasted in the order defined in `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// match taste::<Arch32>(b"\xEB\xFE").nth(0) {
///   Some(loader) => assert_eq!(loader.get_name(), "Windows/32/Raw"),
///   None => panic!("no matching loaders"),
/// };
/// ```
pub fn taste<A: Arch + 'static>(buf: &[u8]) -> impl Iterator<Item=Box<dyn Loader<A>>> {
    default_loaders::<A>()
        .into_iter()
        .filter(move |loader| loader.taste(buf))
}

/// Load the given sample as a 32-bit module using the first matching
///  loader from `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// load::<Arch32>(b"\xEB\xFE")
///   .map(|(loader, module)| {
///     assert_eq!(loader.get_name(),       "Windows/32/Raw");
///     assert_eq!(module.base_address,     0x0);
///     assert_eq!(module.sections[0].name, "raw");
///   })
///   .map_err(|e| panic!(e));
/// ```
pub fn load<A: Arch + 'static>(buf: &[u8]) -> Result<(Box<dyn Loader<A>>, LoadedModule<A>), Error> {
    match taste::<A>(buf).nth(0) {
        Some(loader) => {
            loader.load(buf).map(|module|
                (loader, module)
            )
        },
        None => Err(LoaderError::NotSupported.into()),
    }
}
<|endoftext|>


<commit_before>macro_rules! cstr {
    ($s:expr) => ($s as *const _ as *const ::libc::c_char)
}

macro_rules! real {
    ($f:ident) => ({
        let func = ::hooks::hw::POINTERS.read().unwrap().$f.get();
        func
    })
}

macro_rules! find {
    ($pointers:ident, $handle:ident, $f:ident, $symbol:tt) => ({
        let ptr = $handle.sym($symbol)
            .chain_err(|| concat!("couldn't find ", stringify!($f), "()"))?;
        $pointers.$f.set_from_raw(ptr);
    })
}

macro_rules! command {
    ($name:ident, $callback:expr) => (
        #[allow(non_camel_case_types)]
        struct $name;

        impl $name {
            // This will get called by the engine, in main game thread.
            unsafe extern "C" fn callback() {
                const F: &Fn(::engine::Engine) = &$callback;

                // We know this is the main game thread.
                let engine = ::engine::Engine::new();

                F(engine);
            }

            // This will be called at .so init time. Add this command to the global list.
            #[allow(dead_code)]
            extern "C" fn initialize() {
                #[link_section=".init_array"]
                static INITIALIZE: extern "C" fn() = $name::initialize;

                ::command::COMMANDS.write().unwrap().push(Box::new($name));
            }
        }

        impl ::command::Command for $name {
            fn name(&self) -> &'static [u8] {
                lazy_static! {
                    static ref NAME: ::std::ffi::CString = {
                        ::std::ffi::CString::new(stringify!($name)).unwrap()
                    };
                }

                NAME.as_bytes_with_nul()
            }

            fn callback(&self) -> unsafe extern "C" fn() {
                Self::callback
            }
        }
    )
}

macro_rules! gen_function_impls {
    (@make_impl ($($extern_type:tt)*) ($($arg_name:ident : $arg_type:ident),*)) => (
        impl<R $(, $arg_type)*> Default for Function<$($extern_type)* fn($($arg_type),*) -> R> {
            #[inline(always)]
            fn default() -> Self {
                Function {
                    ptr: Self::default_func as $($extern_type)* fn($($arg_type),*) -> R,
                }
            }
        }

        #[allow(dead_code)]
        impl<R $(, $arg_type)*> Function<$($extern_type)* fn($($arg_type),*) -> R> {
            #[inline(always)]
            pub fn is_default(&self) -> bool {
                self.ptr as *const usize == Self::default_func as *const usize
            }

            #[inline(always)]
            pub unsafe fn call(&self $(, $arg_name : $arg_type)*) -> R {
                (self.ptr)($($arg_name),*)
            }

            #[inline(always)]
            pub fn get(&self) -> $($extern_type)* fn($($arg_type),*) -> R {
                self.ptr
            }

            #[inline(always)]
            pub fn set(&mut self, f: $($extern_type)* fn($($arg_type),*) -> R) {
                self.ptr = f;
            }

            #[inline(always)]
            pub unsafe fn set_from_raw(&mut self, f: *const ::libc::c_void) {
                self.set(*(&f as *const _ as *const _));
            }

            // This should never be called.
            $($extern_type)* fn default_func($(_: $arg_type),*) -> R {
                unreachable!();
            }
        }
    );

    (@gen_impls $($arg_name:ident : $arg_type:ident),*) => (
        gen_function_impls!(@make_impl (unsafe                  ) ($($arg_name : $arg_type),*));
        gen_function_impls!(@make_impl (unsafe extern "C"       ) ($($arg_name : $arg_type),*));
        gen_function_impls!(@make_impl (unsafe extern "system"  ) ($($arg_name : $arg_type),*));
        gen_function_impls!(@make_impl (unsafe extern "fastcall") ($($arg_name : $arg_type),*));
    );

    () => (
        gen_function_impls!(@gen_impls);
    );

    ($first_arg_name:ident : $first_arg_type:ident $(, $arg_name:ident : $arg_type:ident)*) => (
        gen_function_impls!(@gen_impls $first_arg_name : $first_arg_type $(, $arg_name : $arg_type)*);
        gen_function_impls!($($arg_name : $arg_type),*);
    );
}
<commit_msg>Change Command::initialize() to __initialize()<commit_after>macro_rules! cstr {
    ($s:expr) => ($s as *const _ as *const ::libc::c_char)
}

macro_rules! real {
    ($f:ident) => ({
        let func = ::hooks::hw::POINTERS.read().unwrap().$f.get();
        func
    })
}

macro_rules! find {
    ($pointers:ident, $handle:ident, $f:ident, $symbol:tt) => ({
        let ptr = $handle.sym($symbol)
            .chain_err(|| concat!("couldn't find ", stringify!($f), "()"))?;
        $pointers.$f.set_from_raw(ptr);
    })
}

macro_rules! command {
    ($name:ident, $callback:expr) => (
        #[allow(non_camel_case_types)]
        struct $name;

        impl $name {
            // This will get called by the engine, in main game thread.
            unsafe extern "C" fn callback() {
                const F: &Fn(::engine::Engine) = &$callback;

                // We know this is the main game thread.
                let engine = ::engine::Engine::new();

                F(engine);
            }

            // This will be called at .so init time. Add this command to the global list.
            // It shouldn't be called otherwise.
            #[allow(dead_code)]
            extern "C" fn __initialize() {
                #[link_section=".init_array"]
                static INITIALIZE: extern "C" fn() = $name::__initialize;

                ::command::COMMANDS.write().unwrap().push(Box::new($name));
            }
        }

        impl ::command::Command for $name {
            fn name(&self) -> &'static [u8] {
                lazy_static! {
                    static ref NAME: ::std::ffi::CString = {
                        ::std::ffi::CString::new(stringify!($name)).unwrap()
                    };
                }

                NAME.as_bytes_with_nul()
            }

            fn callback(&self) -> unsafe extern "C" fn() {
                Self::callback
            }
        }
    )
}

macro_rules! gen_function_impls {
    (@make_impl ($($extern_type:tt)*) ($($arg_name:ident : $arg_type:ident),*)) => (
        impl<R $(, $arg_type)*> Default for Function<$($extern_type)* fn($($arg_type),*) -> R> {
            #[inline(always)]
            fn default() -> Self {
                Function {
                    ptr: Self::default_func as $($extern_type)* fn($($arg_type),*) -> R,
                }
            }
        }

        #[allow(dead_code)]
        impl<R $(, $arg_type)*> Function<$($extern_type)* fn($($arg_type),*) -> R> {
            #[inline(always)]
            pub fn is_default(&self) -> bool {
                self.ptr as *const usize == Self::default_func as *const usize
            }

            #[inline(always)]
            pub unsafe fn call(&self $(, $arg_name : $arg_type)*) -> R {
                (self.ptr)($($arg_name),*)
            }

            #[inline(always)]
            pub fn get(&self) -> $($extern_type)* fn($($arg_type),*) -> R {
                self.ptr
            }

            #[inline(always)]
            pub fn set(&mut self, f: $($extern_type)* fn($($arg_type),*) -> R) {
                self.ptr = f;
            }

            #[inline(always)]
            pub unsafe fn set_from_raw(&mut self, f: *const ::libc::c_void) {
                self.set(*(&f as *const _ as *const _));
            }

            // This should never be called.
            $($extern_type)* fn default_func($(_: $arg_type),*) -> R {
                unreachable!();
            }
        }
    );

    (@gen_impls $($arg_name:ident : $arg_type:ident),*) => (
        gen_function_impls!(@make_impl (unsafe                  ) ($($arg_name : $arg_type),*));
        gen_function_impls!(@make_impl (unsafe extern "C"       ) ($($arg_name : $arg_type),*));
        gen_function_impls!(@make_impl (unsafe extern "system"  ) ($($arg_name : $arg_type),*));
        gen_function_impls!(@make_impl (unsafe extern "fastcall") ($($arg_name : $arg_type),*));
    );

    () => (
        gen_function_impls!(@gen_impls);
    );

    ($first_arg_name:ident : $first_arg_type:ident $(, $arg_name:ident : $arg_type:ident)*) => (
        gen_function_impls!(@gen_impls $first_arg_name : $first_arg_type $(, $arg_name : $arg_type)*);
        gen_function_impls!($($arg_name : $arg_type),*);
    );
}
<|endoftext|>


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

/// Construct a `serde_json::Value` from a JSON literal.
///
/// ```rust
/// # #[macro_use]
/// # extern crate serde_json;
/// #
/// # fn main() {
/// let value = json!({
///     "code": 200,
///     "success": true,
///     "payload": {
///         "features": [
///             "serde",
///             "json"
///         ]
///     }
/// });
/// # }
/// ```
///
/// Variables or expressions can be interpolated into the JSON literal. Any type
/// interpolated into an array element or object value must implement Serde's
/// `Serialize` trait, while any type interpolated into a object key must
/// implement `Into<String>`. If the `Serialize` implementation of the
/// interpolated type decides to fail, or if the interpolated type contains a
/// map with non-string keys, the `json!` macro will panic.
///
/// ```rust
/// # #[macro_use]
/// # extern crate serde_json;
/// #
/// # fn main() {
/// let code = 200;
/// let features = vec!["serde", "json"];
///
/// let value = json!({
///    "code": code,
///    "success": code == 200,
///    "payload": {
///        features[0]: features[1]
///    }
/// });
/// # }
/// ```
///
/// Trailing commas are allowed inside both arrays and objects.
///
/// ```rust
/// # #[macro_use]
/// # extern crate serde_json;
/// #
/// # fn main() {
/// let value = json!([
///     "notice",
///     "the",
///     "trailing",
///     "comma -->",
/// ]);
/// # }
/// ```
#[macro_export(local_inner_macros)]
macro_rules! json {
    // Hide distracting implementation details from the generated rustdoc.
    ($($json:tt)+) => {
        json_internal!($($json)+)
    };
}

// Rocket relies on this because they export their own `json!` with a different
// doc comment than ours, and various Rust bugs prevent them from calling our
// `json!` from their `json!` so they call `json_internal!` directly. Check with
// @SergioBenitez before making breaking changes to this macro.
//
// Changes are fine as long as `json_internal!` does not call any new helper
// macros and can still be invoked as `json_internal!($($json)+)`.
#[macro_export(local_inner_macros)]
#[doc(hidden)]
macro_rules! json_internal {
    //////////////////////////////////////////////////////////////////////////
    // TT muncher for parsing the inside of an array [...]. Produces a vec![...]
    // of the elements.
    //
    // Must be invoked as: json_internal!(@array [] $($tt)*)
    //////////////////////////////////////////////////////////////////////////

    // Done with trailing comma.
    (@array [$($elems:expr,)*]) => {
        vec![$($elems,)*]
    };

    // Done without trailing comma.
    (@array [$($elems:expr),*]) => {
        vec![$($elems),*]
    };

    // Next element is `null`.
    (@array [$($elems:expr,)*] null $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(null)] $($rest)*)
    };

    // Next element is `true`.
    (@array [$($elems:expr,)*] true $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(true)] $($rest)*)
    };

    // Next element is `false`.
    (@array [$($elems:expr,)*] false $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(false)] $($rest)*)
    };

    // Next element is an array.
    (@array [$($elems:expr,)*] [$($array:tt)*] $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!([$($array)*])] $($rest)*)
    };

    // Next element is a map.
    (@array [$($elems:expr,)*] {$($map:tt)*} $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!({$($map)*})] $($rest)*)
    };

    // Next element is an expression followed by comma.
    (@array [$($elems:expr,)*] $next:expr, $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!($next),] $($rest)*)
    };

    // Last element is an expression with no trailing comma.
    (@array [$($elems:expr,)*] $last:expr) => {
        json_internal!(@array [$($elems,)* json_internal!($last)])
    };

    // Comma after the most recent element.
    (@array [$($elems:expr),*] , $($rest:tt)*) => {
        json_internal!(@array [$($elems,)*] $($rest)*)
    };

    // Unexpected token after most recent element.
    (@array [$($elems:expr),*] $unexpected:tt $($rest:tt)*) => {
        json_unexpected!($unexpected)
    };

    //////////////////////////////////////////////////////////////////////////
    // TT muncher for parsing the inside of an object {...}. Each entry is
    // inserted into the given map variable.
    //
    // Must be invoked as: json_internal!(@object $map () ($($tt)*) ($($tt)*))
    //
    // We require two copies of the input tokens so that we can match on one
    // copy and trigger errors on the other copy.
    //////////////////////////////////////////////////////////////////////////

    // Done.
    (@object $object:ident () () ()) => {};

    // Insert the current entry followed by trailing comma.
    (@object $object:ident [$($key:tt)+] ($value:expr) , $($rest:tt)*) => {
        $object.insert(($($key)+).into(), $value);
        json_internal!(@object $object () ($($rest)*) ($($rest)*));
    };

    // Current entry followed by unexpected token.
    (@object $object:ident [$($key:tt)+] ($value:expr) $unexpected:tt $($rest:tt)*) => {
        json_unexpected!($unexpected);
    };

    // Insert the last entry without trailing comma.
    (@object $object:ident [$($key:tt)+] ($value:expr)) => {
        $object.insert(($($key)+).into(), $value);
    };

    // Next value is `null`.
    (@object $object:ident ($($key:tt)+) (: null $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!(null)) $($rest)*);
    };

    // Next value is `true`.
    (@object $object:ident ($($key:tt)+) (: true $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!(true)) $($rest)*);
    };

    // Next value is `false`.
    (@object $object:ident ($($key:tt)+) (: false $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!(false)) $($rest)*);
    };

    // Next value is an array.
    (@object $object:ident ($($key:tt)+) (: [$($array:tt)*] $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!([$($array)*])) $($rest)*);
    };

    // Next value is a map.
    (@object $object:ident ($($key:tt)+) (: {$($map:tt)*} $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!({$($map)*})) $($rest)*);
    };

    // Next value is an expression followed by comma.
    (@object $object:ident ($($key:tt)+) (: $value:expr , $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!($value)) , $($rest)*);
    };

    // Last value is an expression with no trailing comma.
    (@object $object:ident ($($key:tt)+) (: $value:expr) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!($value)));
    };

    // Missing value for last entry. Trigger a reasonable error message.
    (@object $object:ident ($($key:tt)+) (:) $copy:tt) => {
        // "unexpected end of macro invocation"
        json_internal!();
    };

    // Missing colon and value for last entry. Trigger a reasonable error
    // message.
    (@object $object:ident ($($key:tt)+) () $copy:tt) => {
        // "unexpected end of macro invocation"
        json_internal!();
    };

    // Misplaced colon. Trigger a reasonable error message.
    (@object $object:ident () (: $($rest:tt)*) ($colon:tt $($copy:tt)*)) => {
        // Takes no arguments so "no rules expected the token `:`".
        json_unexpected!($colon);
    };

    // Found a comma inside a key. Trigger a reasonable error message.
    (@object $object:ident ($($key:tt)*) (, $($rest:tt)*) ($comma:tt $($copy:tt)*)) => {
        // Takes no arguments so "no rules expected the token `,`".
        json_unexpected!($comma);
    };

    // Key is fully parenthesized. This avoids clippy double_parens false
    // positives because the parenthesization may be necessary here.
    (@object $object:ident () (($key:expr) : $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object ($key) (: $($rest)*) (: $($rest)*));
    };

    // Munch a token into the current key.
    (@object $object:ident ($($key:tt)*) ($tt:tt $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object ($($key)* $tt) ($($rest)*) ($($rest)*));
    };

    //////////////////////////////////////////////////////////////////////////
    // The main implementation.
    //
    // Must be invoked as: json_internal!($($json)+)
    //////////////////////////////////////////////////////////////////////////

    (null) => {
        $crate::Value::Null
    };

    (true) => {
        $crate::Value::Bool(true)
    };

    (false) => {
        $crate::Value::Bool(false)
    };

    ([]) => {
        $crate::Value::Array(vec![])
    };

    ([ $($tt:tt)+ ]) => {
        $crate::Value::Array(json_internal!(@array [] $($tt)+))
    };

    ({}) => {
        $crate::Value::Object($crate::Map::new())
    };

    ({ $($tt:tt)+ }) => {
        $crate::Value::Object({
            let mut object = $crate::Map::new();
            json_internal!(@object object () ($($tt)+) ($($tt)+));
            object
        })
    };

    // Any Serialize type: numbers, strings, struct literals, variables etc.
    // Must be below every other rule.
    ($other:expr) => {
        $crate::to_value(&$other).unwrap()
    };
}

#[macro_export]
#[doc(hidden)]
macro_rules! json_unexpected {
    () => {};
}
<commit_msg>Revert "Local inner macros"<commit_after>// Copyright 2017 Serde Developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/// Construct a `serde_json::Value` from a JSON literal.
///
/// ```rust
/// # #[macro_use]
/// # extern crate serde_json;
/// #
/// # fn main() {
/// let value = json!({
///     "code": 200,
///     "success": true,
///     "payload": {
///         "features": [
///             "serde",
///             "json"
///         ]
///     }
/// });
/// # }
/// ```
///
/// Variables or expressions can be interpolated into the JSON literal. Any type
/// interpolated into an array element or object value must implement Serde's
/// `Serialize` trait, while any type interpolated into a object key must
/// implement `Into<String>`. If the `Serialize` implementation of the
/// interpolated type decides to fail, or if the interpolated type contains a
/// map with non-string keys, the `json!` macro will panic.
///
/// ```rust
/// # #[macro_use]
/// # extern crate serde_json;
/// #
/// # fn main() {
/// let code = 200;
/// let features = vec!["serde", "json"];
///
/// let value = json!({
///    "code": code,
///    "success": code == 200,
///    "payload": {
///        features[0]: features[1]
///    }
/// });
/// # }
/// ```
///
/// Trailing commas are allowed inside both arrays and objects.
///
/// ```rust
/// # #[macro_use]
/// # extern crate serde_json;
/// #
/// # fn main() {
/// let value = json!([
///     "notice",
///     "the",
///     "trailing",
///     "comma -->",
/// ]);
/// # }
/// ```
#[macro_export]
macro_rules! json {
    // Hide distracting implementation details from the generated rustdoc.
    ($($json:tt)+) => {
        json_internal!($($json)+)
    };
}

// Rocket relies on this because they export their own `json!` with a different
// doc comment than ours, and various Rust bugs prevent them from calling our
// `json!` from their `json!` so they call `json_internal!` directly. Check with
// @SergioBenitez before making breaking changes to this macro.
//
// Changes are fine as long as `json_internal!` does not call any new helper
// macros and can still be invoked as `json_internal!($($json)+)`.
#[macro_export]
#[doc(hidden)]
macro_rules! json_internal {
    //////////////////////////////////////////////////////////////////////////
    // TT muncher for parsing the inside of an array [...]. Produces a vec![...]
    // of the elements.
    //
    // Must be invoked as: json_internal!(@array [] $($tt)*)
    //////////////////////////////////////////////////////////////////////////

    // Done with trailing comma.
    (@array [$($elems:expr,)*]) => {
        vec![$($elems,)*]
    };

    // Done without trailing comma.
    (@array [$($elems:expr),*]) => {
        vec![$($elems),*]
    };

    // Next element is `null`.
    (@array [$($elems:expr,)*] null $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(null)] $($rest)*)
    };

    // Next element is `true`.
    (@array [$($elems:expr,)*] true $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(true)] $($rest)*)
    };

    // Next element is `false`.
    (@array [$($elems:expr,)*] false $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(false)] $($rest)*)
    };

    // Next element is an array.
    (@array [$($elems:expr,)*] [$($array:tt)*] $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!([$($array)*])] $($rest)*)
    };

    // Next element is a map.
    (@array [$($elems:expr,)*] {$($map:tt)*} $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!({$($map)*})] $($rest)*)
    };

    // Next element is an expression followed by comma.
    (@array [$($elems:expr,)*] $next:expr, $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!($next),] $($rest)*)
    };

    // Last element is an expression with no trailing comma.
    (@array [$($elems:expr,)*] $last:expr) => {
        json_internal!(@array [$($elems,)* json_internal!($last)])
    };

    // Comma after the most recent element.
    (@array [$($elems:expr),*] , $($rest:tt)*) => {
        json_internal!(@array [$($elems,)*] $($rest)*)
    };

    // Unexpected token after most recent element.
    (@array [$($elems:expr),*] $unexpected:tt $($rest:tt)*) => {
        json_unexpected!($unexpected)
    };

    //////////////////////////////////////////////////////////////////////////
    // TT muncher for parsing the inside of an object {...}. Each entry is
    // inserted into the given map variable.
    //
    // Must be invoked as: json_internal!(@object $map () ($($tt)*) ($($tt)*))
    //
    // We require two copies of the input tokens so that we can match on one
    // copy and trigger errors on the other copy.
    //////////////////////////////////////////////////////////////////////////

    // Done.
    (@object $object:ident () () ()) => {};

    // Insert the current entry followed by trailing comma.
    (@object $object:ident [$($key:tt)+] ($value:expr) , $($rest:tt)*) => {
        $object.insert(($($key)+).into(), $value);
        json_internal!(@object $object () ($($rest)*) ($($rest)*));
    };

    // Current entry followed by unexpected token.
    (@object $object:ident [$($key:tt)+] ($value:expr) $unexpected:tt $($rest:tt)*) => {
        json_unexpected!($unexpected);
    };

    // Insert the last entry without trailing comma.
    (@object $object:ident [$($key:tt)+] ($value:expr)) => {
        $object.insert(($($key)+).into(), $value);
    };

    // Next value is `null`.
    (@object $object:ident ($($key:tt)+) (: null $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!(null)) $($rest)*);
    };

    // Next value is `true`.
    (@object $object:ident ($($key:tt)+) (: true $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!(true)) $($rest)*);
    };

    // Next value is `false`.
    (@object $object:ident ($($key:tt)+) (: false $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!(false)) $($rest)*);
    };

    // Next value is an array.
    (@object $object:ident ($($key:tt)+) (: [$($array:tt)*] $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!([$($array)*])) $($rest)*);
    };

    // Next value is a map.
    (@object $object:ident ($($key:tt)+) (: {$($map:tt)*} $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!({$($map)*})) $($rest)*);
    };

    // Next value is an expression followed by comma.
    (@object $object:ident ($($key:tt)+) (: $value:expr , $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!($value)) , $($rest)*);
    };

    // Last value is an expression with no trailing comma.
    (@object $object:ident ($($key:tt)+) (: $value:expr) $copy:tt) => {
        json_internal!(@object $object [$($key)+] (json_internal!($value)));
    };

    // Missing value for last entry. Trigger a reasonable error message.
    (@object $object:ident ($($key:tt)+) (:) $copy:tt) => {
        // "unexpected end of macro invocation"
        json_internal!();
    };

    // Missing colon and value for last entry. Trigger a reasonable error
    // message.
    (@object $object:ident ($($key:tt)+) () $copy:tt) => {
        // "unexpected end of macro invocation"
        json_internal!();
    };

    // Misplaced colon. Trigger a reasonable error message.
    (@object $object:ident () (: $($rest:tt)*) ($colon:tt $($copy:tt)*)) => {
        // Takes no arguments so "no rules expected the token `:`".
        json_unexpected!($colon);
    };

    // Found a comma inside a key. Trigger a reasonable error message.
    (@object $object:ident ($($key:tt)*) (, $($rest:tt)*) ($comma:tt $($copy:tt)*)) => {
        // Takes no arguments so "no rules expected the token `,`".
        json_unexpected!($comma);
    };

    // Key is fully parenthesized. This avoids clippy double_parens false
    // positives because the parenthesization may be necessary here.
    (@object $object:ident () (($key:expr) : $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object ($key) (: $($rest)*) (: $($rest)*));
    };

    // Munch a token into the current key.
    (@object $object:ident ($($key:tt)*) ($tt:tt $($rest:tt)*) $copy:tt) => {
        json_internal!(@object $object ($($key)* $tt) ($($rest)*) ($($rest)*));
    };

    //////////////////////////////////////////////////////////////////////////
    // The main implementation.
    //
    // Must be invoked as: json_internal!($($json)+)
    //////////////////////////////////////////////////////////////////////////

    (null) => {
        $crate::Value::Null
    };

    (true) => {
        $crate::Value::Bool(true)
    };

    (false) => {
        $crate::Value::Bool(false)
    };

    ([]) => {
        $crate::Value::Array(vec![])
    };

    ([ $($tt:tt)+ ]) => {
        $crate::Value::Array(json_internal!(@array [] $($tt)+))
    };

    ({}) => {
        $crate::Value::Object($crate::Map::new())
    };

    ({ $($tt:tt)+ }) => {
        $crate::Value::Object({
            let mut object = $crate::Map::new();
            json_internal!(@object object () ($($tt)+) ($($tt)+));
            object
        })
    };

    // Any Serialize type: numbers, strings, struct literals, variables etc.
    // Must be below every other rule.
    ($other:expr) => {
        $crate::to_value(&$other).unwrap()
    };
}

#[macro_export]
#[doc(hidden)]
macro_rules! json_unexpected {
    () => {};
}
<|endoftext|>


<commit_before>#[cfg(perl_multiplicity)]
#[macro_export]
macro_rules! xs_proc {
	($name:ident, $xs:ident, $body:block) => {
		#[allow(dead_code)]
		#[allow(non_snake_case)]
		#[no_mangle]
		pub extern "C" fn $name(perl: &mut $crate::raw::types::PerlInterpreter, cv: &mut $crate::::raw::types::CV) {
			let mut $xs = $crate::XS::new(perl, cv);
			$body
		}
	}
}

#[cfg(not(perl_multiplicity))]
#[macro_export]
macro_rules! xs_proc {
	($name:ident, $xs:ident, $body:block) => {
		#[allow(dead_code)]
		#[allow(non_snake_case)]
		#[no_mangle]
		pub extern "C" fn $name(cv: &mut $crate::raw::types::CV) {
			let mut $xs = $crate::XS::new(cv);
			$body
		}
	}
}

#[macro_export]
macro_rules! XS {
	(package $pkg:ident ($boot:ident) { $( sub $name:ident($xs:ident) $body:block)* }) => {
		$(xs_proc!($name, $xs, $body);)*

		xs_proc!(
			$boot, xs, {
			$({
				let name = [ stringify!($pkg), stringify!($name) ].join("::");
				xs.new_xs(&name, $name, b"Rust code\0");
			})*

			xs.prepush();
			xs.push_long(1);
			xs.putback();
		});
	}
}
<commit_msg>Typo fix.<commit_after>#[cfg(perl_multiplicity)]
#[macro_export]
macro_rules! xs_proc {
	($name:ident, $xs:ident, $body:block) => {
		#[allow(dead_code)]
		#[allow(non_snake_case)]
		#[no_mangle]
		pub extern "C" fn $name(perl: &mut $crate::raw::types::PerlInterpreter, cv: &mut $crate::raw::types::CV) {
			let mut $xs = $crate::XS::new(perl, cv);
			$body
		}
	}
}

#[cfg(not(perl_multiplicity))]
#[macro_export]
macro_rules! xs_proc {
	($name:ident, $xs:ident, $body:block) => {
		#[allow(dead_code)]
		#[allow(non_snake_case)]
		#[no_mangle]
		pub extern "C" fn $name(cv: &mut $crate::raw::types::CV) {
			let mut $xs = $crate::XS::new(cv);
			$body
		}
	}
}

#[macro_export]
macro_rules! XS {
	(package $pkg:ident ($boot:ident) { $( sub $name:ident($xs:ident) $body:block)* }) => {
		$(xs_proc!($name, $xs, $body);)*

		xs_proc!(
			$boot, xs, {
			$({
				let name = [ stringify!($pkg), stringify!($name) ].join("::");
				xs.new_xs(&name, $name, b"Rust code\0");
			})*

			xs.prepush();
			xs.push_long(1);
			xs.putback();
		});
	}
}
<|endoftext|>


<commit_before>use raw::*;
use libc::{c_char};

pub struct Modbus { ctx: *mut modbus_t }

impl Modbus {
    /// Creates a new modbus context with the RTU backend
    ///
    /// # Attributes
    /// * `device`         - Device string e.g. "/dev/ttyUSB0"
    /// * `baud`           - Baud rate
    /// * `parity`         - Parity
    /// * `data_bit`       - Number of data bits
    /// * `stop_bit`       - Number of stop bits
    ///
    /// # Examples
    ///
    /// ```
    /// use libmodbus_rs::modbus::Modbus;
    ///
    /// let modbus = Modbus::new_rtu("/dev/ttyUSB0", 9600, 'N', 8, 1);
    /// ```
    pub fn new_rtu(device: &str, baud: i32, parity: char, data_bit: i32, stop_bit: i32) -> Self {
        Modbus {
            ctx: {
                unsafe {
                    ::raw::modbus_new_rtu(::std::ffi::CString::new(device).unwrap().as_ptr(), baud, parity as c_char, data_bit, stop_bit)
                }
            }
        }
    }

}

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

    #[test]
    fn crate_modbus() {
        let modbus = Modbus::new_rtu("/dev/ttyUSB0", 9600, 'N', 8, 1);
    }

}
<commit_msg>set_slave()<commit_after>use raw::*;
use libc::{c_char};

#[derive(Debug, Eq, PartialEq)]
pub enum ModbusError {
    INVALID_SLAVE_ID,
}

pub struct Modbus { ctx: *mut modbus_t }

impl Modbus {
    /// Creates a new modbus context with the RTU backend
    ///
    /// # Attributes
    /// * `device`         - Device string e.g. "/dev/ttyUSB0"
    /// * `baud`           - Baud rate
    /// * `parity`         - Parity
    /// * `data_bit`       - Number of data bits
    /// * `stop_bit`       - Number of stop bits
    ///
    /// # Examples
    ///
    /// ```
    /// use libmodbus_rs::modbus::Modbus;
    ///
    /// let modbus = Modbus::new_rtu("/dev/ttyUSB0", 9600, 'N', 8, 1);
    /// ```
    pub fn new_rtu(device: &str, baud: i32, parity: char, data_bit: i32, stop_bit: i32) -> Self {
        Modbus {
            ctx: {
                unsafe {
                    ::raw::modbus_new_rtu(::std::ffi::CString::new(device).unwrap().as_ptr(), baud, parity as c_char, data_bit, stop_bit)
                }
            }
        }
    }

    /// Define the slave ID of the remote device to talk in master mode or set the
    /// internal slave ID in slave mode
    ///
    /// # Attributes
    /// * `slave_id`    - New modbus slave id (valid range: `>= 0 && <= 247`)
    ///
    /// # Examples
    /// A `Ok(0)` signals all right, on error a `ModbusError::INVALID_SLAVE_ID` is returned
    ///
    /// ```
    /// use libmodbus_rs::modbus::{Modbus, ModbusError};
    ///
    /// let mut modbus = Modbus::new_rtu("/dev/ttyUSB0", 9600, 'N', 8, 1);
    /// assert_eq!(modbus.set_slave(10), Ok(0));
    /// assert_eq!(modbus.set_slave(255), Err(ModbusError::INVALID_SLAVE_ID));
    /// ```
    pub fn set_slave(&mut self, slave_id: i32) -> Result<i32, ModbusError> {
        unsafe {
            let ret: i32 = ::raw::modbus_set_slave(self.ctx, slave_id);
            if ret == -1 {
                return Err(ModbusError::INVALID_SLAVE_ID);
            } else {
                return Ok(ret);
            }
        }
    }

}

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

    #[test]
    fn crate_modbus() {
        let modbus = Modbus::new_rtu("/dev/ttyUSB0", 9600, 'N', 8, 1);
    }

}
<|endoftext|>