<commit_before>// Copyright 2015-2018 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! This module contains all the types for demuxing DNS oriented streams.

use futures::stream::{Peekable, Stream};
use futures::sync::mpsc::{unbounded, UnboundedReceiver};
use futures::{Async, Future, Poll};

use error::*;
use xfer::{DnsRequest, DnsRequestSender, DnsRequestStreamHandle, DnsResponse, OneshotDnsRequest};

/// This is a generic Exchange implemented over multiplexed DNS connection providers.
///
/// The underlying `DnsRequestSender` is expected to multiplex any I/O connections. DnsExchange assumes that the underlying stream is responsible for this.
#[must_use = "futures do nothing unless polled"]
pub struct DnsExchange<S, R>
where
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    io_stream: S,
    outbound_messages: Peekable<UnboundedReceiver<OneshotDnsRequest<R>>>,
}

impl<S, R> DnsExchange<S, R>
where
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    /// Initializes a TcpStream with an existing tokio_tcp::TcpStream.
    ///
    /// This is intended for use with a TcpListener and Incoming.
    ///
    /// # Arguments
    ///
    /// * `stream` - the established IO stream for communication
    pub fn from_stream(stream: S) -> (Self, DnsRequestStreamHandle<R>) {
        let (message_sender, outbound_messages) = unbounded();
        let message_sender = DnsRequestStreamHandle::<R>::new(message_sender);

        let stream = Self::from_stream_with_receiver(stream, outbound_messages);

        (stream, message_sender)
    }

    /// Wrapps a stream where a sender and receiver have already been established
    pub fn from_stream_with_receiver(
        stream: S,
        receiver: UnboundedReceiver<OneshotDnsRequest<R>>,
    ) -> Self {
        DnsExchange {
            io_stream: stream,
            outbound_messages: receiver.peekable(),
        }
    }

    /// Returns a future, which itself wraps a future which is awaiting connection.
    ///
    /// The connect_future should be lazy.
    pub fn connect<F>(connect_future: F) -> (DnsExchangeConnect<F, S, R>, DnsRequestStreamHandle<R>)
    where
        F: Future<Item = S, Error = ProtoError> + 'static + Send,
    {
        let (message_sender, outbound_messages) = unbounded();
        (
            DnsExchangeConnect::connect(connect_future, outbound_messages),
            DnsRequestStreamHandle::<R>::new(message_sender),
        )
    }
}

impl<S, R> Future for DnsExchange<S, R>
where
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    type Item = ();
    type Error = ProtoError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        // this will not accept incoming data while there is data to send
        //  makes this self throttling.
        loop {
            // poll the underlying stream, to drive it...
            match self.io_stream.poll() {
                // The stream is ready
                Ok(Async::Ready(Some(()))) => (),
                Ok(Async::NotReady) => {
                    if self.io_stream.is_shutdown() {
                        // the io_stream is in a shutdown state, we are only waiting for final results...
                        return Ok(Async::NotReady);
                    }

                    // NotReady and not shutdown, see if there are more messages to send
                    ()
                } // underlying stream is complete.
                Ok(Async::Ready(None)) => {
                    debug!("io_stream is done, shutting down");
                    // TODO: return shutdown error to anything in the stream?

                    return Ok(Async::Ready(()));
                }
                Err(err) => return Err(err),
            }

            // then see if there is more to send
            match self
                .outbound_messages
                .poll()
                .map_err(|()| ProtoError::from("unknown from outbound_message receiver"))?
            {
                // already handled above, here to make sure the poll() pops the next message
                Async::Ready(Some(dns_request)) => {
                    // if there is no peer, this connection should die...
                    let (dns_request, serial_response): (DnsRequest, _) = dns_request.unwrap();

                    info!("sending message via: {}", self.io_stream);

                    match serial_response.send_response(self.io_stream.send_message(dns_request)) {
                        Ok(()) => (),
                        Err(_) => {
                            warn!("failed to associate send_message response to the sender");
                            return Err(
                                "failed to associate send_message response to the sender".into()
                            );
                        }
                    }
                }
                // On not ready, this is our time to return...
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(None) => {
                    debug!("all handles closed, shutting down: {}", self.io_stream);
                    // if there is nothing that can use this connection to send messages, then this is done...
                    self.io_stream.shutdown();

                    // now we'll await the stream to shutdown... see io_stream poll above
                }
            }

            // else we loop to poll on the outbound_messages
        }
    }
}

/// A wrapper for a future DnsExchange connection
pub struct DnsExchangeConnect<F, S, R>(DnsExchangeConnectInner<F, S, R>)
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send;

impl<F, S, R> DnsExchangeConnect<F, S, R>
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    fn connect(
        connect_future: F,
        outbound_messages: UnboundedReceiver<OneshotDnsRequest<R>>,
    ) -> Self {
        DnsExchangeConnect(DnsExchangeConnectInner::Connecting {
            connect_future,
            outbound_messages: Some(outbound_messages),
        })
    }
}

impl<F, S, R> Future for DnsExchangeConnect<F, S, R>
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    type Item = DnsExchange<S, R>;
    type Error = ProtoError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

enum DnsExchangeConnectInner<F, S, R>
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    Connecting {
        connect_future: F,
        outbound_messages: Option<UnboundedReceiver<OneshotDnsRequest<R>>>,
    },
    FailAll {
        error: ProtoError,
        outbound_messages: UnboundedReceiver<OneshotDnsRequest<R>>,
    },
}

impl<F, S, R> Future for DnsExchangeConnectInner<F, S, R>
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    type Item = DnsExchange<S, R>;
    type Error = ProtoError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let next;
            match self {
                DnsExchangeConnectInner::Connecting {
                    ref mut connect_future,
                    ref mut outbound_messages,
                } => {
                    match connect_future.poll() {
                        Ok(Async::Ready(stream)) => {
                            debug!("connection established: {}", stream);
                            return Ok(Async::Ready(DnsExchange::from_stream_with_receiver(
                                stream,
                                outbound_messages
                                    .take()
                                    .expect("cannot poll after complete"),
                            )));
                        }
                        Ok(Async::NotReady) => return Ok(Async::NotReady),
                        Err(error) => {
                            debug!("stream errored while connecting: {:?}", error);
                            next = DnsExchangeConnectInner::FailAll {
                                error,
                                outbound_messages: outbound_messages
                                    .take()
                                    .expect("cannot poll after complete"),
                            }
                        }
                    };
                }
                DnsExchangeConnectInner::FailAll {
                    error,
                    ref mut outbound_messages,
                } => {
                    while let Some(outbound_message) = match outbound_messages.poll() {
                        Ok(Async::Ready(opt)) => opt,
                        Ok(Async::NotReady) => return Ok(Async::NotReady),
                        Err(_) => None,
                    } {
                        let response = S::error_response(error.clone());
                        // ignoring errors... best effort send...
                        outbound_message.unwrap().1.send_response(response).ok();
                    }

                    return Err(error.clone());
                }
            }

            *self = next;
        }
    }
}
<commit_msg>Changed log message to debug.<commit_after>// Copyright 2015-2018 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! This module contains all the types for demuxing DNS oriented streams.

use futures::stream::{Peekable, Stream};
use futures::sync::mpsc::{unbounded, UnboundedReceiver};
use futures::{Async, Future, Poll};

use error::*;
use xfer::{DnsRequest, DnsRequestSender, DnsRequestStreamHandle, DnsResponse, OneshotDnsRequest};

/// This is a generic Exchange implemented over multiplexed DNS connection providers.
///
/// The underlying `DnsRequestSender` is expected to multiplex any I/O connections. DnsExchange assumes that the underlying stream is responsible for this.
#[must_use = "futures do nothing unless polled"]
pub struct DnsExchange<S, R>
where
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    io_stream: S,
    outbound_messages: Peekable<UnboundedReceiver<OneshotDnsRequest<R>>>,
}

impl<S, R> DnsExchange<S, R>
where
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    /// Initializes a TcpStream with an existing tokio_tcp::TcpStream.
    ///
    /// This is intended for use with a TcpListener and Incoming.
    ///
    /// # Arguments
    ///
    /// * `stream` - the established IO stream for communication
    pub fn from_stream(stream: S) -> (Self, DnsRequestStreamHandle<R>) {
        let (message_sender, outbound_messages) = unbounded();
        let message_sender = DnsRequestStreamHandle::<R>::new(message_sender);

        let stream = Self::from_stream_with_receiver(stream, outbound_messages);

        (stream, message_sender)
    }

    /// Wrapps a stream where a sender and receiver have already been established
    pub fn from_stream_with_receiver(
        stream: S,
        receiver: UnboundedReceiver<OneshotDnsRequest<R>>,
    ) -> Self {
        DnsExchange {
            io_stream: stream,
            outbound_messages: receiver.peekable(),
        }
    }

    /// Returns a future, which itself wraps a future which is awaiting connection.
    ///
    /// The connect_future should be lazy.
    pub fn connect<F>(connect_future: F) -> (DnsExchangeConnect<F, S, R>, DnsRequestStreamHandle<R>)
    where
        F: Future<Item = S, Error = ProtoError> + 'static + Send,
    {
        let (message_sender, outbound_messages) = unbounded();
        (
            DnsExchangeConnect::connect(connect_future, outbound_messages),
            DnsRequestStreamHandle::<R>::new(message_sender),
        )
    }
}

impl<S, R> Future for DnsExchange<S, R>
where
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    type Item = ();
    type Error = ProtoError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        // this will not accept incoming data while there is data to send
        //  makes this self throttling.
        loop {
            // poll the underlying stream, to drive it...
            match self.io_stream.poll() {
                // The stream is ready
                Ok(Async::Ready(Some(()))) => (),
                Ok(Async::NotReady) => {
                    if self.io_stream.is_shutdown() {
                        // the io_stream is in a shutdown state, we are only waiting for final results...
                        return Ok(Async::NotReady);
                    }

                    // NotReady and not shutdown, see if there are more messages to send
                    ()
                } // underlying stream is complete.
                Ok(Async::Ready(None)) => {
                    debug!("io_stream is done, shutting down");
                    // TODO: return shutdown error to anything in the stream?

                    return Ok(Async::Ready(()));
                }
                Err(err) => return Err(err),
            }

            // then see if there is more to send
            match self
                .outbound_messages
                .poll()
                .map_err(|()| ProtoError::from("unknown from outbound_message receiver"))?
            {
                // already handled above, here to make sure the poll() pops the next message
                Async::Ready(Some(dns_request)) => {
                    // if there is no peer, this connection should die...
                    let (dns_request, serial_response): (DnsRequest, _) = dns_request.unwrap();

                    debug!("sending message via: {}", self.io_stream);

                    match serial_response.send_response(self.io_stream.send_message(dns_request)) {
                        Ok(()) => (),
                        Err(_) => {
                            warn!("failed to associate send_message response to the sender");
                            return Err(
                                "failed to associate send_message response to the sender".into()
                            );
                        }
                    }
                }
                // On not ready, this is our time to return...
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(None) => {
                    debug!("all handles closed, shutting down: {}", self.io_stream);
                    // if there is nothing that can use this connection to send messages, then this is done...
                    self.io_stream.shutdown();

                    // now we'll await the stream to shutdown... see io_stream poll above
                }
            }

            // else we loop to poll on the outbound_messages
        }
    }
}

/// A wrapper for a future DnsExchange connection
pub struct DnsExchangeConnect<F, S, R>(DnsExchangeConnectInner<F, S, R>)
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send;

impl<F, S, R> DnsExchangeConnect<F, S, R>
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    fn connect(
        connect_future: F,
        outbound_messages: UnboundedReceiver<OneshotDnsRequest<R>>,
    ) -> Self {
        DnsExchangeConnect(DnsExchangeConnectInner::Connecting {
            connect_future,
            outbound_messages: Some(outbound_messages),
        })
    }
}

impl<F, S, R> Future for DnsExchangeConnect<F, S, R>
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    type Item = DnsExchange<S, R>;
    type Error = ProtoError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

enum DnsExchangeConnectInner<F, S, R>
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    Connecting {
        connect_future: F,
        outbound_messages: Option<UnboundedReceiver<OneshotDnsRequest<R>>>,
    },
    FailAll {
        error: ProtoError,
        outbound_messages: UnboundedReceiver<OneshotDnsRequest<R>>,
    },
}

impl<F, S, R> Future for DnsExchangeConnectInner<F, S, R>
where
    F: Future<Item = S, Error = ProtoError> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R>,
    R: Future<Item = DnsResponse, Error = ProtoError> + 'static + Send,
{
    type Item = DnsExchange<S, R>;
    type Error = ProtoError;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let next;
            match self {
                DnsExchangeConnectInner::Connecting {
                    ref mut connect_future,
                    ref mut outbound_messages,
                } => {
                    match connect_future.poll() {
                        Ok(Async::Ready(stream)) => {
                            debug!("connection established: {}", stream);
                            return Ok(Async::Ready(DnsExchange::from_stream_with_receiver(
                                stream,
                                outbound_messages
                                    .take()
                                    .expect("cannot poll after complete"),
                            )));
                        }
                        Ok(Async::NotReady) => return Ok(Async::NotReady),
                        Err(error) => {
                            debug!("stream errored while connecting: {:?}", error);
                            next = DnsExchangeConnectInner::FailAll {
                                error,
                                outbound_messages: outbound_messages
                                    .take()
                                    .expect("cannot poll after complete"),
                            }
                        }
                    };
                }
                DnsExchangeConnectInner::FailAll {
                    error,
                    ref mut outbound_messages,
                } => {
                    while let Some(outbound_message) = match outbound_messages.poll() {
                        Ok(Async::Ready(opt)) => opt,
                        Ok(Async::NotReady) => return Ok(Async::NotReady),
                        Err(_) => None,
                    } {
                        let response = S::error_response(error.clone());
                        // ignoring errors... best effort send...
                        outbound_message.unwrap().1.send_response(response).ok();
                    }

                    return Err(error.clone());
                }
            }

            *self = next;
        }
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "nightly-testing", plugin(clippy))]
#![cfg_attr(feature = "nightly-testing", feature(plugin))]
#![cfg_attr(feature = "nightly-testing", allow(too_many_arguments))]
#![cfg_attr(feature = "nightly-testing", allow(used_underscore_binding))]
#![cfg_attr(feature = "nightly-testing", allow(useless_let_if_seq))]
#![cfg_attr(not(feature = "with-syntex"), feature(rustc_private, plugin))]
#![cfg_attr(not(feature = "with-syntex"), plugin(quasi_macros))]

extern crate aster;
extern crate quasi;

#[cfg(feature = "with-syntex")]
extern crate syntex;

#[cfg(feature = "with-syntex")]
#[macro_use]
extern crate syntex_syntax as syntax;

#[cfg(not(feature = "with-syntex"))]
#[macro_use]
extern crate syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate rustc_plugin;

#[cfg(not(feature = "with-syntex"))]
use syntax::feature_gate::AttributeType;

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/lib.rs"));

#[cfg(not(feature = "with-syntex"))]
include!("lib.rs.in");

#[cfg(feature = "with-syntex")]
pub fn register(reg: &mut syntex::Registry) {
    use syntax::{ast, fold};

    /// Strip the serde attributes from the crate.
    #[cfg(feature = "with-syntex")]
    fn strip_attributes(krate: ast::Crate) -> ast::Crate {
        /// Helper folder that strips the serde attributes after the extensions have been expanded.
        struct StripAttributeFolder;

        impl fold::Folder for StripAttributeFolder {
            fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
                match attr.node.value.node {
                    ast::MetaItemKind::List(ref n, _) if n == &"serde" => { return None; }
                    _ => {}
                }

                Some(attr)
            }

            fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
                fold::noop_fold_mac(mac, self)
            }
        }

        fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
    }

    reg.add_attr("feature(custom_derive)");
    reg.add_attr("feature(custom_attribute)");

    reg.add_decorator("derive_Serialize", ser::expand_derive_serialize);
    reg.add_decorator("derive_Deserialize", de::expand_derive_deserialize);

    reg.add_post_expansion_pass(strip_attributes);
}

#[cfg(not(feature = "with-syntex"))]
pub fn register(reg: &mut rustc_plugin::Registry) {
    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Serialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(ser::expand_derive_serialize)));

    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Deserialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(de::expand_derive_deserialize)));

    reg.register_attribute("serde".to_owned(), AttributeType::Normal);
}
<commit_msg>Re-enable clippy lint "useless_let_if_seq"<commit_after>#![cfg_attr(feature = "nightly-testing", plugin(clippy))]
#![cfg_attr(feature = "nightly-testing", feature(plugin))]
#![cfg_attr(feature = "nightly-testing", allow(too_many_arguments))]
#![cfg_attr(feature = "nightly-testing", allow(used_underscore_binding))]
#![cfg_attr(not(feature = "with-syntex"), feature(rustc_private, plugin))]
#![cfg_attr(not(feature = "with-syntex"), plugin(quasi_macros))]

extern crate aster;
extern crate quasi;

#[cfg(feature = "with-syntex")]
extern crate syntex;

#[cfg(feature = "with-syntex")]
#[macro_use]
extern crate syntex_syntax as syntax;

#[cfg(not(feature = "with-syntex"))]
#[macro_use]
extern crate syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate rustc_plugin;

#[cfg(not(feature = "with-syntex"))]
use syntax::feature_gate::AttributeType;

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/lib.rs"));

#[cfg(not(feature = "with-syntex"))]
include!("lib.rs.in");

#[cfg(feature = "with-syntex")]
pub fn register(reg: &mut syntex::Registry) {
    use syntax::{ast, fold};

    /// Strip the serde attributes from the crate.
    #[cfg(feature = "with-syntex")]
    fn strip_attributes(krate: ast::Crate) -> ast::Crate {
        /// Helper folder that strips the serde attributes after the extensions have been expanded.
        struct StripAttributeFolder;

        impl fold::Folder for StripAttributeFolder {
            fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
                match attr.node.value.node {
                    ast::MetaItemKind::List(ref n, _) if n == &"serde" => { return None; }
                    _ => {}
                }

                Some(attr)
            }

            fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
                fold::noop_fold_mac(mac, self)
            }
        }

        fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
    }

    reg.add_attr("feature(custom_derive)");
    reg.add_attr("feature(custom_attribute)");

    reg.add_decorator("derive_Serialize", ser::expand_derive_serialize);
    reg.add_decorator("derive_Deserialize", de::expand_derive_deserialize);

    reg.add_post_expansion_pass(strip_attributes);
}

#[cfg(not(feature = "with-syntex"))]
pub fn register(reg: &mut rustc_plugin::Registry) {
    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Serialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(ser::expand_derive_serialize)));

    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Deserialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(de::expand_derive_deserialize)));

    reg.register_attribute("serde".to_owned(), AttributeType::Normal);
}
<|endoftext|>


<commit_before>extern crate bindgen;

use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;

#[allow(dead_code)]
fn parse_ld_archive(ar: &Path) -> Vec<String> {
    let mut f = File::open(ar).unwrap();
    let mut content = String::new();
    f.read_to_string(&mut content).unwrap();
    if "GROUP" == &content[0..5] {
        println!("Found group");
        let open_idx = content.find("(").unwrap_or_else(|| content.len());
        let remove_open = content[open_idx + 1..].trim();
        let end_idx = remove_open.find(")").unwrap_or_else(|| remove_open.len());
        let remaining = remove_open[..end_idx].trim();
        println!("Remaining is {}", remaining);
        remaining
            .split_whitespace()
            .map(|s| {
                let end = s.len() - 2;
                String::from(&s[3..end])
            })
            .collect()
    } else {
        panic!("Could not find a group");
    }
}

#[allow(dead_code)]
fn write_external_link(libs: &Vec<String>) {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest = Path::new(&out_dir).join("linkage.rs");
    let mut f = File::create(&dest).unwrap();
    for l in libs {
        let link_str = format!("#[link(name=\"{}\", kind=\"static\")]", l);
        let overall = link_str + "\nextern \"C\" {}\n";
        f.write_all(&overall.into_bytes()).unwrap();
    }
}

/// Cargo runs main in this file to get some additional settings (e.g., LD_LIBRARY_PATH). It reads the printed output
/// looking for certain variables, see [here](http://doc.crates.io/build-script.html) for documentation.
fn main() {
    // Get the directory where we are building.
    let cargo_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
    let dpdk_build = Path::new(&cargo_dir)
        .parent()
        .unwrap()
        .join("3rdparty")
        .join("dpdk")
        .join("build");

    let dpdk_libs = dpdk_build.clone().join("lib");
    let native_path = Path::new(&cargo_dir)
        .parent()
        .unwrap()
        .join("target")
        .join("native");
    //println!("DPDK {:?}", dpdk_libs.to_str());
    // Use DPDK directory as -L
    println!(
        "cargo:rustc-link-search=native={}",
        dpdk_libs.to_str().unwrap()
    );
    if dpdk_libs.join("libdpdk.so").exists() {
        println!("cargo:rustc-link-lib=dpdk");
    }
    println!(
        "cargo:rustc-link-search=native={}",
        native_path.to_str().unwrap()
    );
    let header_path = Path::new(&cargo_dir)
        .join("src")
        .join("native_include")
        .join("dpdk-headers.h");
    let dpdk_include_path = dpdk_build.clone().join("include");
    println!("Header path {:?}", header_path.to_str());
    let bindings = bindgen::Builder::default()
        .header(header_path.to_str().unwrap())
        .rust_target(bindgen::RustTarget::Nightly)
        .clang_args(vec!["-I", dpdk_include_path.to_str().unwrap()].iter())
        .generate()
        .expect("Unable to generate DPDK bindings");
    let out_dir = env::var("OUT_DIR").unwrap();
    let dpdk_bindings = Path::new(&out_dir).join("dpdk_bindings.rs");
    bindings
        .write_to_file(dpdk_bindings)
        .expect("Could not write bindings");
}
<commit_msg>Fix #90 Remove max_align_t from generate binding<commit_after>extern crate bindgen;

use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;

#[allow(dead_code)]
fn parse_ld_archive(ar: &Path) -> Vec<String> {
    let mut f = File::open(ar).unwrap();
    let mut content = String::new();
    f.read_to_string(&mut content).unwrap();
    if "GROUP" == &content[0..5] {
        println!("Found group");
        let open_idx = content.find("(").unwrap_or_else(|| content.len());
        let remove_open = content[open_idx + 1..].trim();
        let end_idx = remove_open.find(")").unwrap_or_else(|| remove_open.len());
        let remaining = remove_open[..end_idx].trim();
        println!("Remaining is {}", remaining);
        remaining
            .split_whitespace()
            .map(|s| {
                let end = s.len() - 2;
                String::from(&s[3..end])
            })
            .collect()
    } else {
        panic!("Could not find a group");
    }
}

#[allow(dead_code)]
fn write_external_link(libs: &Vec<String>) {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest = Path::new(&out_dir).join("linkage.rs");
    let mut f = File::create(&dest).unwrap();
    for l in libs {
        let link_str = format!("#[link(name=\"{}\", kind=\"static\")]", l);
        let overall = link_str + "\nextern \"C\" {}\n";
        f.write_all(&overall.into_bytes()).unwrap();
    }
}

/// Cargo runs main in this file to get some additional settings (e.g., LD_LIBRARY_PATH). It reads the printed output
/// looking for certain variables, see [here](http://doc.crates.io/build-script.html) for documentation.
fn main() {
    // Get the directory where we are building.
    let cargo_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
    let dpdk_build = Path::new(&cargo_dir)
        .parent()
        .unwrap()
        .join("3rdparty")
        .join("dpdk")
        .join("build");

    let dpdk_libs = dpdk_build.clone().join("lib");
    let native_path = Path::new(&cargo_dir)
        .parent()
        .unwrap()
        .join("target")
        .join("native");
    //println!("DPDK {:?}", dpdk_libs.to_str());
    // Use DPDK directory as -L
    println!(
        "cargo:rustc-link-search=native={}",
        dpdk_libs.to_str().unwrap()
    );
    if dpdk_libs.join("libdpdk.so").exists() {
        println!("cargo:rustc-link-lib=dpdk");
    }
    println!(
        "cargo:rustc-link-search=native={}",
        native_path.to_str().unwrap()
    );
    let header_path = Path::new(&cargo_dir)
        .join("src")
        .join("native_include")
        .join("dpdk-headers.h");
    let dpdk_include_path = dpdk_build.clone().join("include");
    println!("Header path {:?}", header_path.to_str());
    let bindings = bindgen::Builder::default()
        .header(header_path.to_str().unwrap())
        .rust_target(bindgen::RustTarget::Nightly)
        .clang_args(vec!["-I", dpdk_include_path.to_str().unwrap()].iter())
        .blacklist_type("max_align_t") // https://github.com/servo/rust-bindgen/issues/550
        .generate()
        .expect("Unable to generate DPDK bindings");
    let out_dir = env::var("OUT_DIR").unwrap();
    let dpdk_bindings = Path::new(&out_dir).join("dpdk_bindings.rs");
    bindings
        .write_to_file(dpdk_bindings)
        .expect("Could not write bindings");
}
<|endoftext|>


<commit_before>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::io;
use std::os::unix::io::AsRawFd;
use std::os::unix::net::UnixListener;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::{anyhow, bail, Context};
use argh::FromArgs;
use base::{
    error, get_max_open_files, warn, AsRawDescriptors, Event, RawDescriptor, Tube,
    UnlinkUnixListener,
};
use cros_async::{EventAsync, Executor};
use data_model::{DataInit, Le32};
use fuse::Server;
use futures::future::{AbortHandle, Abortable};
use hypervisor::ProtectionType;
use minijail::{self, Minijail};
use sync::Mutex;
use vm_memory::GuestMemory;
use vmm_vhost::message::{VhostUserProtocolFeatures, VhostUserVirtioFeatures};

use crate::virtio;
use crate::virtio::copy_config;
use crate::virtio::fs::passthrough::PassthroughFs;
use crate::virtio::fs::{process_fs_queue, virtio_fs_config, FS_MAX_TAG_LEN};
use crate::virtio::vhost::user::device::handler::{
    DeviceRequestHandler, Doorbell, VhostUserBackend,
};

async fn handle_fs_queue(
    mut queue: virtio::Queue,
    mem: GuestMemory,
    doorbell: Arc<Mutex<Doorbell>>,
    kick_evt: EventAsync,
    server: Arc<fuse::Server<PassthroughFs>>,
    tube: Arc<Mutex<Tube>>,
) {
    // Slot is always going to be 0 because we do not support DAX
    let slot: u32 = 0;

    loop {
        if let Err(e) = kick_evt.next_val().await {
            error!("Failed to read kick event for fs queue: {}", e);
            break;
        }
        if let Err(e) = process_fs_queue(&mem, &doorbell, &mut queue, &server, &tube, slot) {
            error!("Process FS queue failed: {}", e);
            break;
        }
    }
}

fn default_uidmap() -> String {
    let euid = unsafe { libc::geteuid() };
    format!("{} {} 1", euid, euid)
}

fn default_gidmap() -> String {
    let egid = unsafe { libc::getegid() };
    format!("{} {} 1", egid, egid)
}

fn jail_and_fork(
    mut keep_rds: Vec<RawDescriptor>,
    dir_path: PathBuf,
    uid_map: Option<String>,
    gid_map: Option<String>,
) -> anyhow::Result<i32> {
    // Create new minijail sandbox
    let mut j = Minijail::new()?;

    j.namespace_pids();
    j.namespace_user();
    j.namespace_user_disable_setgroups();
    j.uidmap(&uid_map.unwrap_or_else(default_uidmap))?;
    j.gidmap(&gid_map.unwrap_or_else(default_gidmap))?;
    j.run_as_init();

    j.namespace_vfs();
    j.namespace_net();
    j.no_new_privs();

    // Only pivot_root if we are not re-using the current root directory.
    if dir_path != Path::new("/") {
        // It's safe to call `namespace_vfs` multiple times.
        j.namespace_vfs();
        j.enter_pivot_root(&dir_path)?;
    }
    j.set_remount_mode(libc::MS_SLAVE);

    let limit = get_max_open_files().context("failed to get max open files")?;
    j.set_rlimit(libc::RLIMIT_NOFILE as i32, limit, limit)?;

    // Make sure there are no duplicates in keep_rds
    keep_rds.dedup();

    // fork on the jail here
    let pid = unsafe { j.fork(Some(&keep_rds))? };

    if pid > 0 {
        unsafe { libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) };
    }

    if pid < 0 {
        bail!("Fork error! {}", std::io::Error::last_os_error());
    }

    Ok(pid)
}

struct FsBackend {
    ex: Executor,
    server: Arc<fuse::Server<PassthroughFs>>,
    tag: [u8; FS_MAX_TAG_LEN],
    avail_features: u64,
    acked_features: u64,
    acked_protocol_features: VhostUserProtocolFeatures,
    workers: [Option<AbortHandle>; Self::MAX_QUEUE_NUM],
    keep_rds: Vec<RawDescriptor>,
}

impl FsBackend {
    pub fn new(ex: &Executor, tag: &str) -> anyhow::Result<Self> {
        if tag.len() > FS_MAX_TAG_LEN {
            bail!(
                "fs tag is too long: {} (max supported: {})",
                tag.len(),
                FS_MAX_TAG_LEN
            );
        }
        let mut fs_tag = [0u8; FS_MAX_TAG_LEN];
        fs_tag[..tag.len()].copy_from_slice(tag.as_bytes());

        let avail_features = virtio::base_features(ProtectionType::Unprotected)
            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();

        // Use default passthroughfs config
        let fs = PassthroughFs::new(Default::default())?;

        let mut keep_rds: Vec<RawDescriptor> = [0, 1, 2].to_vec();
        keep_rds.append(&mut fs.keep_rds());

        let ex = ex.clone();
        keep_rds.extend(ex.as_raw_descriptors());

        let server = Arc::new(Server::new(fs));

        Ok(FsBackend {
            ex,
            server,
            tag: fs_tag,
            avail_features,
            acked_features: 0,
            acked_protocol_features: VhostUserProtocolFeatures::empty(),
            workers: Default::default(),
            keep_rds,
        })
    }
}

impl VhostUserBackend for FsBackend {
    const MAX_QUEUE_NUM: usize = 2; /* worker queue and high priority queue */
    const MAX_VRING_LEN: u16 = 1024;

    type Error = anyhow::Error;

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, value: u64) -> anyhow::Result<()> {
        let unrequested_features = value & !self.avail_features;
        if unrequested_features != 0 {
            bail!("invalid features are given: {:#x}", unrequested_features);
        }

        self.acked_features |= value;

        Ok(())
    }

    fn acked_features(&self) -> u64 {
        self.acked_features
    }

    fn protocol_features(&self) -> VhostUserProtocolFeatures {
        VhostUserProtocolFeatures::CONFIG | VhostUserProtocolFeatures::MQ
    }

    fn ack_protocol_features(&mut self, features: u64) -> anyhow::Result<()> {
        let features = VhostUserProtocolFeatures::from_bits(features)
            .ok_or_else(|| anyhow!("invalid protocol features are given: {:#x}", features))?;
        let supported = self.protocol_features();
        self.acked_protocol_features = features & supported;
        Ok(())
    }

    fn acked_protocol_features(&self) -> u64 {
        self.acked_protocol_features.bits()
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let config = virtio_fs_config {
            tag: self.tag,
            num_request_queues: Le32::from(1),
        };
        copy_config(data, 0, config.as_slice(), offset);
    }

    fn reset(&mut self) {
        for handle in self.workers.iter_mut().filter_map(Option::take) {
            handle.abort();
        }
    }

    fn start_queue(
        &mut self,
        idx: usize,
        mut queue: virtio::Queue,
        mem: GuestMemory,
        doorbell: Arc<Mutex<Doorbell>>,
        kick_evt: Event,
    ) -> anyhow::Result<()> {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            warn!("Starting new queue handler without stopping old handler");
            handle.abort();
        }

        // Enable any virtqueue features that were negotiated (like VIRTIO_RING_F_EVENT_IDX).
        queue.ack_features(self.acked_features);

        let kick_evt = EventAsync::new(kick_evt, &self.ex)
            .context("failed to create EventAsync for kick_evt")?;
        let (handle, registration) = AbortHandle::new_pair();
        let (_, fs_device_tube) = Tube::pair()?;

        self.ex
            .spawn_local(Abortable::new(
                handle_fs_queue(
                    queue,
                    mem,
                    doorbell,
                    kick_evt,
                    self.server.clone(),
                    Arc::new(Mutex::new(fs_device_tube)),
                ),
                registration,
            ))
            .detach();

        self.workers[idx] = Some(handle);
        Ok(())
    }

    fn stop_queue(&mut self, idx: usize) {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            handle.abort();
        }
    }
}

#[derive(FromArgs)]
#[argh(description = "")]
struct Options {
    #[argh(option, description = "path to a socket", arg_name = "PATH")]
    socket: String,
    #[argh(option, description = "the virtio-fs tag", arg_name = "TAG")]
    tag: String,
    #[argh(option, description = "path to a directory to share", arg_name = "DIR")]
    shared_dir: PathBuf,
    #[argh(option, description = "uid map to use", arg_name = "UIDMAP")]
    uid_map: Option<String>,
    #[argh(option, description = "gid map to use", arg_name = "GIDMAP")]
    gid_map: Option<String>,
}

/// Starts a vhost-user fs device.
/// Returns an error if the given `args` is invalid or the device fails to run.
pub fn run_fs_device(program_name: &str, args: &[&str]) -> anyhow::Result<()> {
    let opts = match Options::from_args(&[program_name], args) {
        Ok(opts) => opts,
        Err(e) => {
            if e.status.is_err() {
                bail!(e.output);
            } else {
                println!("{}", e.output);
            }
            return Ok(());
        }
    };

    base::syslog::init().context("Failed to initialize syslog")?;

    let ex = Executor::new().context("Failed to create executor")?;
    let fs_device = FsBackend::new(&ex, &opts.tag)?;

    // Create and bind unix socket
    let listener = UnixListener::bind(opts.socket).map(UnlinkUnixListener)?;
    let mut keep_rds = fs_device.keep_rds.clone();
    keep_rds.push(listener.as_raw_fd());
    base::syslog::push_descriptors(&mut keep_rds);

    let handler = DeviceRequestHandler::new(fs_device);

    let pid = jail_and_fork(keep_rds, opts.shared_dir, opts.uid_map, opts.gid_map)?;

    // Parent, nothing to do but wait and then exit
    if pid != 0 {
        unsafe { libc::waitpid(pid, std::ptr::null_mut(), 0) };
        return Ok(());
    }

    // We need to set the no setuid fixup secure bit so that we don't drop capabilities when
    // changing the thread uid/gid. Without this, creating new entries can fail in some corner
    // cases.
    const SECBIT_NO_SETUID_FIXUP: i32 = 1 << 2;
    // TODO(crbug.com/1199487): Remove this once libc provides the wrapper for all targets.
    #[cfg(target_os = "linux")]
    {
        // Safe because this doesn't modify any memory and we check the return value.
        let mut securebits = unsafe { libc::prctl(libc::PR_GET_SECUREBITS) };
        if securebits < 0 {
            bail!(io::Error::last_os_error());
        }
        securebits |= SECBIT_NO_SETUID_FIXUP;
        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe { libc::prctl(libc::PR_SET_SECUREBITS, securebits) };
        if ret < 0 {
            bail!(io::Error::last_os_error());
        }
    }

    // run_until() returns an Result<Result<..>> which the ? operator lets us flatten.
    ex.run_until(handler.run_with_listener(listener, &ex))?
}
<commit_msg>devices: vhost-user: fs: support VVU<commit_after>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::io;
use std::os::unix::io::AsRawFd;
use std::os::unix::net::UnixListener;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::{anyhow, bail, Context};
use argh::FromArgs;
use base::{
    error, get_max_open_files, warn, AsRawDescriptors, Event, RawDescriptor, Tube,
    UnlinkUnixListener,
};
use cros_async::{EventAsync, Executor};
use data_model::{DataInit, Le32};
use fuse::Server;
use futures::future::{AbortHandle, Abortable};
use hypervisor::ProtectionType;
use minijail::{self, Minijail};
use sync::Mutex;
use vm_memory::GuestMemory;
use vmm_vhost::message::{VhostUserProtocolFeatures, VhostUserVirtioFeatures};

use crate::virtio;
use crate::virtio::copy_config;
use crate::virtio::fs::passthrough::PassthroughFs;
use crate::virtio::fs::{process_fs_queue, virtio_fs_config, FS_MAX_TAG_LEN};
use crate::virtio::vhost::user::device::handler::{
    DeviceRequestHandler, Doorbell, VhostUserBackend,
};
use crate::virtio::vhost::user::device::vvu::pci::VvuPciDevice;

async fn handle_fs_queue(
    mut queue: virtio::Queue,
    mem: GuestMemory,
    doorbell: Arc<Mutex<Doorbell>>,
    kick_evt: EventAsync,
    server: Arc<fuse::Server<PassthroughFs>>,
    tube: Arc<Mutex<Tube>>,
) {
    // Slot is always going to be 0 because we do not support DAX
    let slot: u32 = 0;

    loop {
        if let Err(e) = kick_evt.next_val().await {
            error!("Failed to read kick event for fs queue: {}", e);
            break;
        }
        if let Err(e) = process_fs_queue(&mem, &doorbell, &mut queue, &server, &tube, slot) {
            error!("Process FS queue failed: {}", e);
            break;
        }
    }
}

fn default_uidmap() -> String {
    let euid = unsafe { libc::geteuid() };
    format!("{} {} 1", euid, euid)
}

fn default_gidmap() -> String {
    let egid = unsafe { libc::getegid() };
    format!("{} {} 1", egid, egid)
}

fn jail_and_fork(
    mut keep_rds: Vec<RawDescriptor>,
    dir_path: PathBuf,
    uid_map: Option<String>,
    gid_map: Option<String>,
) -> anyhow::Result<i32> {
    // Create new minijail sandbox
    let mut j = Minijail::new()?;

    j.namespace_pids();
    j.namespace_user();
    j.namespace_user_disable_setgroups();
    j.uidmap(&uid_map.unwrap_or_else(default_uidmap))?;
    j.gidmap(&gid_map.unwrap_or_else(default_gidmap))?;
    j.run_as_init();

    j.namespace_vfs();
    j.namespace_net();
    j.no_new_privs();

    // Only pivot_root if we are not re-using the current root directory.
    if dir_path != Path::new("/") {
        // It's safe to call `namespace_vfs` multiple times.
        j.namespace_vfs();
        j.enter_pivot_root(&dir_path)?;
    }
    j.set_remount_mode(libc::MS_SLAVE);

    let limit = get_max_open_files().context("failed to get max open files")?;
    j.set_rlimit(libc::RLIMIT_NOFILE as i32, limit, limit)?;
    // vvu locks around 512k memory. Just give 1M.
    j.set_rlimit(libc::RLIMIT_MEMLOCK as i32, 1 << 20, 1 << 20)?;

    // Make sure there are no duplicates in keep_rds
    keep_rds.dedup();

    // fork on the jail here
    let pid = unsafe { j.fork(Some(&keep_rds))? };

    if pid > 0 {
        unsafe { libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM) };
    }

    if pid < 0 {
        bail!("Fork error! {}", std::io::Error::last_os_error());
    }

    Ok(pid)
}

struct FsBackend {
    ex: Executor,
    server: Arc<fuse::Server<PassthroughFs>>,
    tag: [u8; FS_MAX_TAG_LEN],
    avail_features: u64,
    acked_features: u64,
    acked_protocol_features: VhostUserProtocolFeatures,
    workers: [Option<AbortHandle>; Self::MAX_QUEUE_NUM],
    keep_rds: Vec<RawDescriptor>,
}

impl FsBackend {
    pub fn new(ex: &Executor, tag: &str) -> anyhow::Result<Self> {
        if tag.len() > FS_MAX_TAG_LEN {
            bail!(
                "fs tag is too long: {} (max supported: {})",
                tag.len(),
                FS_MAX_TAG_LEN
            );
        }
        let mut fs_tag = [0u8; FS_MAX_TAG_LEN];
        fs_tag[..tag.len()].copy_from_slice(tag.as_bytes());

        let avail_features = virtio::base_features(ProtectionType::Unprotected)
            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();

        // Use default passthroughfs config
        let fs = PassthroughFs::new(Default::default())?;

        let mut keep_rds: Vec<RawDescriptor> = [0, 1, 2].to_vec();
        keep_rds.append(&mut fs.keep_rds());

        let ex = ex.clone();
        keep_rds.extend(ex.as_raw_descriptors());

        let server = Arc::new(Server::new(fs));

        Ok(FsBackend {
            ex,
            server,
            tag: fs_tag,
            avail_features,
            acked_features: 0,
            acked_protocol_features: VhostUserProtocolFeatures::empty(),
            workers: Default::default(),
            keep_rds,
        })
    }
}

impl VhostUserBackend for FsBackend {
    const MAX_QUEUE_NUM: usize = 2; /* worker queue and high priority queue */
    const MAX_VRING_LEN: u16 = 1024;

    type Error = anyhow::Error;

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, value: u64) -> anyhow::Result<()> {
        let unrequested_features = value & !self.avail_features;
        if unrequested_features != 0 {
            bail!("invalid features are given: {:#x}", unrequested_features);
        }

        self.acked_features |= value;

        Ok(())
    }

    fn acked_features(&self) -> u64 {
        self.acked_features
    }

    fn protocol_features(&self) -> VhostUserProtocolFeatures {
        VhostUserProtocolFeatures::CONFIG | VhostUserProtocolFeatures::MQ
    }

    fn ack_protocol_features(&mut self, features: u64) -> anyhow::Result<()> {
        let features = VhostUserProtocolFeatures::from_bits(features)
            .ok_or_else(|| anyhow!("invalid protocol features are given: {:#x}", features))?;
        let supported = self.protocol_features();
        self.acked_protocol_features = features & supported;
        Ok(())
    }

    fn acked_protocol_features(&self) -> u64 {
        self.acked_protocol_features.bits()
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        let config = virtio_fs_config {
            tag: self.tag,
            num_request_queues: Le32::from(1),
        };
        copy_config(data, 0, config.as_slice(), offset);
    }

    fn reset(&mut self) {
        for handle in self.workers.iter_mut().filter_map(Option::take) {
            handle.abort();
        }
    }

    fn start_queue(
        &mut self,
        idx: usize,
        mut queue: virtio::Queue,
        mem: GuestMemory,
        doorbell: Arc<Mutex<Doorbell>>,
        kick_evt: Event,
    ) -> anyhow::Result<()> {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            warn!("Starting new queue handler without stopping old handler");
            handle.abort();
        }

        // Enable any virtqueue features that were negotiated (like VIRTIO_RING_F_EVENT_IDX).
        queue.ack_features(self.acked_features);

        let kick_evt = EventAsync::new(kick_evt, &self.ex)
            .context("failed to create EventAsync for kick_evt")?;
        let (handle, registration) = AbortHandle::new_pair();
        let (_, fs_device_tube) = Tube::pair()?;

        self.ex
            .spawn_local(Abortable::new(
                handle_fs_queue(
                    queue,
                    mem,
                    doorbell,
                    kick_evt,
                    self.server.clone(),
                    Arc::new(Mutex::new(fs_device_tube)),
                ),
                registration,
            ))
            .detach();

        self.workers[idx] = Some(handle);
        Ok(())
    }

    fn stop_queue(&mut self, idx: usize) {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            handle.abort();
        }
    }
}

#[derive(FromArgs)]
#[argh(description = "")]
struct Options {
    #[argh(option, description = "path to a vhost-user socket", arg_name = "PATH")]
    socket: Option<String>,
    #[argh(
        option,
        description = "VFIO-PCI device name (e.g. '0000:00:07.0')",
        arg_name = "STRING"
    )]
    vfio: Option<String>,
    #[argh(option, description = "the virtio-fs tag", arg_name = "TAG")]
    tag: String,
    #[argh(option, description = "path to a directory to share", arg_name = "DIR")]
    shared_dir: PathBuf,
    #[argh(option, description = "uid map to use", arg_name = "UIDMAP")]
    uid_map: Option<String>,
    #[argh(option, description = "gid map to use", arg_name = "GIDMAP")]
    gid_map: Option<String>,
}

/// Starts a vhost-user fs device.
/// Returns an error if the given `args` is invalid or the device fails to run.
pub fn run_fs_device(program_name: &str, args: &[&str]) -> anyhow::Result<()> {
    let opts = match Options::from_args(&[program_name], args) {
        Ok(opts) => opts,
        Err(e) => {
            if e.status.is_err() {
                bail!(e.output);
            } else {
                println!("{}", e.output);
            }
            return Ok(());
        }
    };

    base::syslog::init().context("Failed to initialize syslog")?;

    let ex = Executor::new().context("Failed to create executor")?;
    let fs_device = FsBackend::new(&ex, &opts.tag)?;

    let mut keep_rds = fs_device.keep_rds.clone();
    let listener = match opts.socket {
        None => None,
        Some(socket) => {
            // Create and bind unix socket
            let l = UnixListener::bind(socket).map(UnlinkUnixListener)?;
            keep_rds.push(l.as_raw_fd());
            Some(l)
        }
    };
    base::syslog::push_descriptors(&mut keep_rds);

    let handler = DeviceRequestHandler::new(fs_device);

    let pid = jail_and_fork(keep_rds, opts.shared_dir, opts.uid_map, opts.gid_map)?;

    // Parent, nothing to do but wait and then exit
    if pid != 0 {
        unsafe { libc::waitpid(pid, std::ptr::null_mut(), 0) };
        return Ok(());
    }

    // We need to set the no setuid fixup secure bit so that we don't drop capabilities when
    // changing the thread uid/gid. Without this, creating new entries can fail in some corner
    // cases.
    const SECBIT_NO_SETUID_FIXUP: i32 = 1 << 2;
    // TODO(crbug.com/1199487): Remove this once libc provides the wrapper for all targets.
    #[cfg(target_os = "linux")]
    {
        // Safe because this doesn't modify any memory and we check the return value.
        let mut securebits = unsafe { libc::prctl(libc::PR_GET_SECUREBITS) };
        if securebits < 0 {
            bail!(io::Error::last_os_error());
        }
        securebits |= SECBIT_NO_SETUID_FIXUP;
        // Safe because this doesn't modify any memory and we check the return value.
        let ret = unsafe { libc::prctl(libc::PR_SET_SECUREBITS, securebits) };
        if ret < 0 {
            bail!(io::Error::last_os_error());
        }
    }

    let res = match (listener, opts.vfio) {
        (Some(l), None) => ex.run_until(handler.run_with_listener(l, &ex))?,
        (None, Some(vfio)) => {
            let device = VvuPciDevice::new(&vfio, FsBackend::MAX_QUEUE_NUM)?;
            ex.run_until(handler.run_vvu(device, &ex))?
        }
        _ => Err(anyhow!("exactly one of `--socket` or `--vfio` is required")),
    };
    res
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate grin_api as api;
extern crate grin_chain as chain;
extern crate grin_core as core;
extern crate grin_p2p as p2p;
extern crate grin_servers as servers;
extern crate grin_util as util;
extern crate grin_wallet as wallet;

extern crate bufstream;
extern crate serde_json;

mod framework;

use bufstream::BufStream;
use serde_json::Value;
use std::io::prelude::{BufRead, Write};
use std::net::TcpStream;

use std::{thread, time};

use core::global::{self, ChainTypes};

use framework::{config, stratum_config};

// Create a grin server, and a stratum server.
// Simulate a few JSONRpc requests and verify the results.
// Validate disconnected workers
// Validate broadcasting new jobs
#[test]
fn basic_stratum_server() {
	util::init_test_logger();
	global::set_mining_mode(ChainTypes::AutomatedTesting);

	let test_name_dir = "stratum_server";
	framework::clean_all_output(test_name_dir);

	// Create a server
	let s = servers::Server::new(config(4000, test_name_dir, 0)).unwrap();

	// Get mining config with stratumserver enabled
	let mut stratum_cfg = stratum_config();
	stratum_cfg.burn_reward = true;
	stratum_cfg.attempt_time_per_block = 999;
	stratum_cfg.enable_stratum_server = Some(true);
	stratum_cfg.stratum_server_addr = Some(String::from("127.0.0.1:11101"));

	// Start stratum server
	s.start_stratum_server(stratum_cfg);

	// Wait for stratum server to start and
	// Verify stratum server accepts connections
	loop {
		if let Ok(_stream) = TcpStream::connect("127.0.0.1:11101") {
			break;
		} else {
			thread::sleep(time::Duration::from_millis(500));
		}
		// As this stream falls out of scope it will be disconnected
	}

	// Create a few new worker connections
	let mut workers = vec![];
	for _n in 0..5 {
		let w = TcpStream::connect("127.0.0.1:11101").unwrap();
		w.set_nonblocking(true)
			.expect("Failed to set TcpStream to non-blocking");
		let stream = BufStream::new(w);
		workers.push(stream);
	}
	assert!(workers.len() == 5);

	// Simulate a worker lost connection
	workers.remove(4);

	// Swallow the genesis block
	thread::sleep(time::Duration::from_secs(5)); // Wait for the server to broadcast
	let mut response = String::new();
	for n in 0..workers.len() {
		let _result = workers[n].read_line(&mut response);
	}

	// Verify a few stratum JSONRpc commands
	// getjobtemplate - expected block template result
	let mut response = String::new();
	let job_req = "{\"id\": \"Stratum\", \"jsonrpc\": \"2.0\", \"method\": \"getjobtemplate\"}\n";
	workers[2].write(job_req.as_bytes()).unwrap();
	workers[2].flush().unwrap();
	thread::sleep(time::Duration::from_secs(1)); // Wait for the server to reply
	match workers[2].read_line(&mut response) {
		Ok(_) => {
			let r: Value = serde_json::from_str(&response).unwrap();
			assert_eq!(r["error"], serde_json::Value::Null);
			assert_ne!(r["result"], serde_json::Value::Null);
		}
		Err(_e) => {
			assert!(false);
		}
	}

	// keepalive - expected "ok" result
	let mut response = String::new();
	let job_req = "{\"id\":\"3\",\"jsonrpc\":\"2.0\",\"method\":\"keepalive\"}\n";
	let ok_resp = "{\"id\":\"3\",\"jsonrpc\":\"2.0\",\"method\":\"keepalive\",\"result\":\"ok\",\"error\":null}\n";
	workers[2].write(job_req.as_bytes()).unwrap();
	workers[2].flush().unwrap();
	thread::sleep(time::Duration::from_secs(1)); // Wait for the server to reply
	let _st = workers[2].read_line(&mut response);
	assert_eq!(response.as_str(), ok_resp);

	// "doesnotexist" - error expected
	let mut response = String::new();
	let job_req = "{\"id\":\"4\",\"jsonrpc\":\"2.0\",\"method\":\"doesnotexist\"}\n";
	let ok_resp = "{\"id\":\"4\",\"jsonrpc\":\"2.0\",\"method\":\"doesnotexist\",\"result\":null,\"error\":{\"code\":-32601,\"message\":\"Method not found\"}}\n";
	workers[3].write(job_req.as_bytes()).unwrap();
	workers[3].flush().unwrap();
	thread::sleep(time::Duration::from_secs(1)); // Wait for the server to reply
	let _st = workers[3].read_line(&mut response);
	assert_eq!(response.as_str(), ok_resp);

	// Verify stratum server and worker stats
	let stats = s.get_server_stats().unwrap();
	assert_eq!(stats.stratum_stats.block_height, 1); // just 1 genesis block
	assert_eq!(stats.stratum_stats.num_workers, 4); // 5 - 1 = 4
	assert_eq!(stats.stratum_stats.worker_stats[5].is_connected, false); // worker was removed
	assert_eq!(stats.stratum_stats.worker_stats[1].is_connected, true);

	// Start mining blocks
	s.start_test_miner(None);

	// Simulate a worker lost connection
	workers.remove(1);

	// Verify blocks are being broadcast to workers
	let expected = String::from("job");
	thread::sleep(time::Duration::from_secs(3)); // Wait for a few mined blocks
	let mut jobtemplate = String::new();
	let _st = workers[2].read_line(&mut jobtemplate);
	let job_template: Value = serde_json::from_str(&jobtemplate).unwrap();
	assert_eq!(job_template["method"], expected);

	// Verify stratum server and worker stats
	let stats = s.get_server_stats().unwrap();
	assert_eq!(stats.stratum_stats.num_workers, 3); // 5 - 2 = 3
	assert_eq!(stats.stratum_stats.worker_stats[2].is_connected, false); // worker was removed
	assert_ne!(stats.stratum_stats.block_height, 1);
}
<commit_msg>test: avoid infinite waiting on basic_stratum_server test (#1634)<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate grin_api as api;
extern crate grin_chain as chain;
extern crate grin_core as core;
extern crate grin_p2p as p2p;
extern crate grin_servers as servers;
extern crate grin_util as util;
extern crate grin_wallet as wallet;

extern crate bufstream;
extern crate serde_json;
#[macro_use]
extern crate slog;

mod framework;

use bufstream::BufStream;
use serde_json::Value;
use std::io::prelude::{BufRead, Write};
use std::net::TcpStream;

use std::process;
use std::{thread, time};

use core::global::{self, ChainTypes};
use util::LOGGER;

use framework::{config, stratum_config};

// Create a grin server, and a stratum server.
// Simulate a few JSONRpc requests and verify the results.
// Validate disconnected workers
// Validate broadcasting new jobs
#[test]
fn basic_stratum_server() {
	util::init_test_logger();
	global::set_mining_mode(ChainTypes::AutomatedTesting);

	let test_name_dir = "stratum_server";
	framework::clean_all_output(test_name_dir);

	// Create a server
	let s = servers::Server::new(config(4000, test_name_dir, 0)).unwrap();

	// Get mining config with stratumserver enabled
	let mut stratum_cfg = stratum_config();
	stratum_cfg.burn_reward = true;
	stratum_cfg.attempt_time_per_block = 999;
	stratum_cfg.enable_stratum_server = Some(true);
	stratum_cfg.stratum_server_addr = Some(String::from("127.0.0.1:11101"));

	// Start stratum server
	s.start_stratum_server(stratum_cfg);

	// Wait for stratum server to start and
	// Verify stratum server accepts connections
	loop {
		if let Ok(_stream) = TcpStream::connect("127.0.0.1:11101") {
			break;
		} else {
			thread::sleep(time::Duration::from_millis(500));
		}
		// As this stream falls out of scope it will be disconnected
	}
	info!(LOGGER, "stratum server connected");

	// Create a few new worker connections
	let mut workers = vec![];
	for _n in 0..5 {
		let w = TcpStream::connect("127.0.0.1:11101").unwrap();
		w.set_nonblocking(true)
			.expect("Failed to set TcpStream to non-blocking");
		let stream = BufStream::new(w);
		workers.push(stream);
	}
	assert!(workers.len() == 5);
	info!(LOGGER, "workers length verification ok");

	// Simulate a worker lost connection
	workers.remove(4);

	// Swallow the genesis block
	thread::sleep(time::Duration::from_secs(5)); // Wait for the server to broadcast
	let mut response = String::new();
	for n in 0..workers.len() {
		let _result = workers[n].read_line(&mut response);
	}

	// Verify a few stratum JSONRpc commands
	// getjobtemplate - expected block template result
	let mut response = String::new();
	let job_req = "{\"id\": \"Stratum\", \"jsonrpc\": \"2.0\", \"method\": \"getjobtemplate\"}\n";
	workers[2].write(job_req.as_bytes()).unwrap();
	workers[2].flush().unwrap();
	thread::sleep(time::Duration::from_secs(1)); // Wait for the server to reply
	match workers[2].read_line(&mut response) {
		Ok(_) => {
			let r: Value = serde_json::from_str(&response).unwrap();
			assert_eq!(r["error"], serde_json::Value::Null);
			assert_ne!(r["result"], serde_json::Value::Null);
		}
		Err(_e) => {
			assert!(false);
		}
	}
	info!(LOGGER, "a few stratum JSONRpc commands verification ok");

	// keepalive - expected "ok" result
	let mut response = String::new();
	let job_req = "{\"id\":\"3\",\"jsonrpc\":\"2.0\",\"method\":\"keepalive\"}\n";
	let ok_resp = "{\"id\":\"3\",\"jsonrpc\":\"2.0\",\"method\":\"keepalive\",\"result\":\"ok\",\"error\":null}\n";
	workers[2].write(job_req.as_bytes()).unwrap();
	workers[2].flush().unwrap();
	thread::sleep(time::Duration::from_secs(1)); // Wait for the server to reply
	let _st = workers[2].read_line(&mut response);
	assert_eq!(response.as_str(), ok_resp);
	info!(LOGGER, "keepalive test ok");

	// "doesnotexist" - error expected
	let mut response = String::new();
	let job_req = "{\"id\":\"4\",\"jsonrpc\":\"2.0\",\"method\":\"doesnotexist\"}\n";
	let ok_resp = "{\"id\":\"4\",\"jsonrpc\":\"2.0\",\"method\":\"doesnotexist\",\"result\":null,\"error\":{\"code\":-32601,\"message\":\"Method not found\"}}\n";
	workers[3].write(job_req.as_bytes()).unwrap();
	workers[3].flush().unwrap();
	thread::sleep(time::Duration::from_secs(1)); // Wait for the server to reply
	let _st = workers[3].read_line(&mut response);
	assert_eq!(response.as_str(), ok_resp);
	info!(LOGGER, "worker doesnotexist test ok");

	// Verify stratum server and worker stats
	let stats = s.get_server_stats().unwrap();
	assert_eq!(stats.stratum_stats.block_height, 1); // just 1 genesis block
	assert_eq!(stats.stratum_stats.num_workers, 4); // 5 - 1 = 4
	assert_eq!(stats.stratum_stats.worker_stats[5].is_connected, false); // worker was removed
	assert_eq!(stats.stratum_stats.worker_stats[1].is_connected, true);
	info!(LOGGER, "stratum server and worker stats verification ok");

	// Start mining blocks
	s.start_test_miner(None);
	info!(LOGGER, "test miner started");

	// This test is supposed to complete in 3 seconds,
	// so let's set a timeout on 10s to avoid infinite waiting happened in Travis-CI.
	let handler = thread::spawn(|| {
		thread::sleep(time::Duration::from_secs(10));
		error!(LOGGER, "basic_stratum_server test fail on timeout!");
		thread::sleep(time::Duration::from_millis(100));
		process::exit(1);
	});

	// Simulate a worker lost connection
	workers.remove(1);

	// Verify blocks are being broadcast to workers
	let expected = String::from("job");
	thread::sleep(time::Duration::from_secs(3)); // Wait for a few mined blocks
	let mut jobtemplate = String::new();
	let _st = workers[2].read_line(&mut jobtemplate);
	let job_template: Value = serde_json::from_str(&jobtemplate).unwrap();
	assert_eq!(job_template["method"], expected);
	info!(LOGGER, "blocks broadcasting to workers test ok");

	// Verify stratum server and worker stats
	let stats = s.get_server_stats().unwrap();
	assert_eq!(stats.stratum_stats.num_workers, 3); // 5 - 2 = 3
	assert_eq!(stats.stratum_stats.worker_stats[2].is_connected, false); // worker was removed
	assert_ne!(stats.stratum_stats.block_height, 1);
	info!(LOGGER, "basic_stratum_server test done and ok.");
}
<|endoftext|>


<commit_before>use core::cmp::{Eq, Ord};
use iter::BaseIter;

pub trait PersistentMap<K: Copy Eq Ord, V: Copy> {
  pure fn get(k: K) -> Option<V>;
  pure fn put(k: K, v: V) -> self;
  pure fn delete(k: K) -> self;
}

enum RBColor {
  Red,
  Black
}

pub enum RBMap<K: Copy Eq Ord, V: Copy> {
  //   color    left         k  value       right
  Tree(RBColor, @RBMap<K,V>, K, Option<V>, @RBMap<K,V>),
  Leaf
}

pure fn RBMap<K: Copy Eq Ord, V: Copy>(key: K, value: V) -> @RBMap<K, V> {
  @Tree(Red, @Leaf, key, Some(value), @Leaf)
}

priv impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> {
  pure fn modify(k: K, new_value: Option<V>) -> @RBMap<K, V> {
    match self {
      @Leaf => @Tree(Red, self, k, copy new_value, self),
      @Tree(color, left, key, original_value, right) => {
        if (k.lt(&key)) {
          self.balanceLeft(color, left.modify(k, new_value), key, original_value, right)
        } else if (k == key) {
          @Tree(color, left, k, copy new_value, right)
        } else {
          self.balanceRight(color, left, key, original_value, right.modify(k, new_value))
        }
      }
    }
  }

  pure fn balanceLeft(c: RBColor, l: @RBMap<K,V>, k: K, v: Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black,@Tree(Red,@Tree(Red,a,xK,xV,b),yK,yV,c),zK,zV,d) =>
        @Tree(Red,@Tree(Black,a,xK,xV,b),yK,yV,@Tree(Black,c,zK,zV,d)),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }

  pure fn balanceRight(c: RBColor, l: @RBMap<K,V>, k: K, v: Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black, @Tree(Red, t1, k1, v1, t2), k0, v0, @Tree(Red, t3, k2, v2, t4)) =>
        @Tree(Red, @Tree(Black, t1, k1, v1, t2), k0, v0, @Tree(Black, t3, k2, v2, t4)),
      (c, l, k, v, @Tree(Red, t1, k1, v1, t2)) =>
        @Tree(c, @Tree(Red, l, k, v, t1), k1, v1, t2),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }
}

/**
 * A purely functional Left-Leaning Red-Black Tree.
 */
impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> : PersistentMap<K, V> {
  pure fn get(k: K) -> Option<V> {
    match self {
      @Leaf => None,
      @Tree(_, left, key, maybe_value, right) => {
        if (k.lt(&key)) {
          left.get(k)
        } else if (k.gt(&key)) {
          right.get(k)
        } else {
          copy maybe_value
        }
      }
    }
  }

  pure fn put(k: K, new_value: V) -> @RBMap<K, V> {
    self.modify(k, Some(new_value))
  }

  pure fn delete(k: K) -> @RBMap<K, V> {
    self.modify(k, None)
  }
}

impl<K: Copy Eq Ord, V: Copy> RBMap<K, V>: BaseIter<(&K, &V)> {
  pure fn size_hint(&self) -> Option<uint> {
    None
  }

  pure fn each(&self, f: fn(&(&self/K, &self/V)) -> bool) {
    match *self {
      Leaf => (),
      Tree(_, left, key, maybe_value, right) => {
        left.each(f);
        match maybe_value {
          Some(value) => f(&(&key, &value)),
          None => false,
        };
        right.each(f);
      }
    }
  }
}

#[test]
fn test_rb_tree() {
  let v1 = RBMap("stevej", 150);
  let v2 = v1.put("thatstacy", 187);

  assert(v2.get("stevej") == Some(150));
  assert(v2.get("thatstacy") == Some(187));

  let v3 = v2.delete("stevej");
  assert(v3.get("stevej") == None);
  assert(v3.get("thatstacy") == Some(187));

  let v4 = v2.put("jeremy", 16);

  let v5 = v4.put("ev", 20);
  let v6 = v5.put("zhanna", 51573);

  assert(v6.get("jeremy") == Some(16));
  assert(v6.get("stevej") == Some(150));
  assert(v6.get("thatstacy") == Some(187));
  assert(v6.get("ev") == Some(20));
  assert(v6.get("zhana") == Some(51573));
}

#[test]
fn test_base_iter_each() {
  let v1 = RBMap(1, 0);
  let v2 = v1.put(4, 0);
  let v3 = v2.put(3, 0);
  let v4 = v3.put(5, 0);
  let v5 = v4.put(2, 0);

  let n = @mut 1;

  fn t(n: @mut int, kv: &(&int, &int)) -> bool{
    match kv {
      &(k, _) => {
        assert (*n == *k);
        *n += 1;
      }
    }
    false
  }

  v5.each(|z| t(n, z));
}
<commit_msg>Correct borrowing errors.<commit_after>use core::cmp::{Eq, Ord};
use iter::BaseIter;

pub trait PersistentMap<K: Copy Eq Ord, V: Copy> {
  pure fn get(k: K) -> Option<V>;
  pure fn put(k: K, v: V) -> self;
  pure fn delete(k: K) -> self;
}

enum RBColor {
  Red,
  Black
}

pub enum RBMap<K: Copy Eq Ord, V: Copy> {
  //   color    left         k  value       right
  Tree(RBColor, @RBMap<K,V>, K, Option<V>, @RBMap<K,V>),
  Leaf
}

pure fn RBMap<K: Copy Eq Ord, V: Copy>(key: K, value: V) -> @RBMap<K, V> {
  @Tree(Red, @Leaf, key, Some(value), @Leaf)
}

priv impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> {
  pure fn modify(k: K, new_value: Option<V>) -> @RBMap<K, V> {
    match self {
      @Leaf => @Tree(Red, self, k, copy new_value, self),
      @Tree(color, left, key, original_value, right) => {
        if (k.lt(&key)) {
          self.balanceLeft(color, left.modify(k, new_value), key, original_value, right)
        } else if (k == key) {
          @Tree(color, left, k, copy new_value, right)
        } else {
          self.balanceRight(color, left, key, original_value, right.modify(k, new_value))
        }
      }
    }
  }

  pure fn balanceLeft(c: RBColor, l: @RBMap<K,V>, k: K, v: Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black,@Tree(Red,@Tree(Red,a,xK,xV,b),yK,yV,c),zK,zV,d) =>
        @Tree(Red,@Tree(Black,a,xK,xV,b),yK,yV,@Tree(Black,c,zK,zV,d)),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }

  pure fn balanceRight(c: RBColor, l: @RBMap<K,V>, k: K, v: Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black, @Tree(Red, t1, k1, v1, t2), k0, v0, @Tree(Red, t3, k2, v2, t4)) =>
        @Tree(Red, @Tree(Black, t1, k1, v1, t2), k0, v0, @Tree(Black, t3, k2, v2, t4)),
      (c, l, k, v, @Tree(Red, t1, k1, v1, t2)) =>
        @Tree(c, @Tree(Red, l, k, v, t1), k1, v1, t2),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }
}

/**
 * A purely functional Left-Leaning Red-Black Tree.
 */
impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> : PersistentMap<K, V> {
  pure fn get(k: K) -> Option<V> {
    match self {
      @Leaf => None,
      @Tree(_, left, key, maybe_value, right) => {
        if (k.lt(&key)) {
          left.get(k)
        } else if (k.gt(&key)) {
          right.get(k)
        } else {
          copy maybe_value
        }
      }
    }
  }

  pure fn put(k: K, new_value: V) -> @RBMap<K, V> {
    self.modify(k, Some(new_value))
  }

  pure fn delete(k: K) -> @RBMap<K, V> {
    self.modify(k, None)
  }
}

impl<K: Copy Eq Ord, V: Copy> RBMap<K, V>: BaseIter<(&K, &V)> {
  pure fn size_hint(&self) -> Option<uint> {
    None
  }

  pure fn each(&self, f: fn(&(&self/K, &self/V)) -> bool) {
    match *self {
      Leaf => (),
      Tree(_, ref left, ref key, ref maybe_value, ref right) => {
        let left: &self/@RBMap<K,V> = left;
        let key: &self/K = key;
        let maybe_value: &self/Option<V> = maybe_value;
        let right: &self/@RBMap<K,V> = right;
        left.each(f);
        match *maybe_value {
            Some(ref value) => {
                let value: &self/V = value;
                f(&(key, value));
            }
            None => ()
        };
        right.each(f);
      }
    }
  }
}

#[test]
fn test_rb_tree() {
  let v1 = RBMap("stevej", 150);
  let v2 = v1.put("thatstacy", 187);

  assert(v2.get("stevej") == Some(150));
  assert(v2.get("thatstacy") == Some(187));

  let v3 = v2.delete("stevej");
  assert(v3.get("stevej") == None);
  assert(v3.get("thatstacy") == Some(187));

  let v4 = v2.put("jeremy", 16);

  let v5 = v4.put("ev", 20);
  let v6 = v5.put("zhanna", 51573);

  assert(v6.get("jeremy") == Some(16));
  assert(v6.get("stevej") == Some(150));
  assert(v6.get("thatstacy") == Some(187));
  assert(v6.get("ev") == Some(20));
  assert(v6.get("zhana") == Some(51573));
}

#[test]
fn test_base_iter_each() {
  let v1 = RBMap(1, 0);
  let v2 = v1.put(4, 0);
  let v3 = v2.put(3, 0);
  let v4 = v3.put(5, 0);
  let v5 = v4.put(2, 0);

  let n = @mut 1;

  fn t(n: @mut int, kv: &(&int, &int)) -> bool{
    match kv {
      &(k, _) => {
        assert (*n == *k);
        *n += 1;
      }
    }
    false
  }

  v5.each(|z| t(n, z));
}
<|endoftext|>


<commit_before>//! This is the actual Bot module. For ergonomic reasons there is a RcBot which composes the real bot
//! as an underlying field. You should always use RcBot.

use objects;
use error::Error;

use std::str;
use std::io;
use std::time::Duration;
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::{RefCell, Cell};
use std::sync::{Arc, Mutex};

use curl::easy::{Easy, List, Form, InfoType};
use tokio_curl::Session;
use tokio_core::reactor::{Handle, Core, Interval};
use serde_json;
use serde_json::value::Value;
use futures::{Future, IntoFuture, Stream, stream};
use futures::sync::mpsc;
use futures::sync::mpsc::UnboundedSender;

/// A clonable, single threaded bot
///
/// The outer API gets implemented on RcBot
#[derive(Clone)]
pub struct RcBot {
    pub inner: Rc<Bot>,
}

impl RcBot {
    pub fn new<'a>(handle: Handle, key: &str) -> impl Future<Item = RcBot, Error = Error> + 'a {
        use functions::FunctionGetMe;
        let bot = RcBot { inner: Rc::new(Bot::new(handle, key)) };
        bot.get_me().send().map(|(mut bot, me)| {
            Rc::get_mut(&mut bot.inner).as_mut().unwrap().id = me.id;
            Rc::get_mut(&mut bot.inner).as_mut().unwrap().username = me.username.unwrap();
            bot
        })
    }
}

/// The main bot structure
pub struct Bot {
    pub key: String,
    pub handle: Handle,
    pub last_id: Cell<u32>,
    pub update_interval: Cell<u64>,
    pub handlers: RefCell<HashMap<String, UnboundedSender<(RcBot, objects::Message)>>>,
    pub session: Session,
    pub username: String,
    pub id: i64,
}

impl Bot {
    pub fn new(handle: Handle, key: &str) -> Bot {
        Bot {
            handle: handle.clone(),
            key: key.into(),
            last_id: Cell::new(0),
            update_interval: Cell::new(1000),
            handlers: RefCell::new(HashMap::new()),
            session: Session::new(handle.clone()),
            username: String::new(),
            id: 0,
        }
    }

    /// Creates a new request and adds a JSON message to it. The returned Future contains a the
    /// reply as a string.  This method should be used if no file is added because a JSON msg is
    /// always compacter than a formdata one.
    pub fn fetch_json<'a>(&self, func: &str, msg: &str) -> impl Future<Item = String, Error = Error> + 'a {
        println!("Send JSON: {}", msg);

        let mut header = List::new();
        header.append("Content-Type: application/json").unwrap();

        let mut a = Easy::new();
        a.http_headers(header).unwrap();
        a.post_fields_copy(msg.as_bytes()).unwrap();
        a.post(true).unwrap();

        self._fetch(func, a)
    }

    /// Creates a new request with some byte content (e.g. a file). The method properties have to be
    /// in the formdata setup and cannot be sent as JSON.
    pub fn fetch_formdata<'a, T>(&self,
                                 func: &str,
                                 msg: Value,
                                 mut file: T,
                                 kind: &str,
                                 file_name: &str)
                                 -> impl Future<Item = String, Error = Error> + 'a
        where T: io::Read
    {
        let mut content = Vec::new();

        let mut a = Easy::new();
        let mut form = Form::new();

        // try to read the byte content to a vector
        let _ = file.read_to_end(&mut content).unwrap();

        // add properties
        for (key, val) in msg.as_object().unwrap().iter() {
            form.part(key).contents(format!("{:?}", val).as_bytes()).add().unwrap();
        }

        // add the file
        form.part(kind).buffer(file_name, content).content_type("application/octet-stream").add().unwrap();

        // create a http post request
        a.post(true).unwrap();
        a.httppost(form).unwrap();

        self._fetch(func, a)
    }

    /// calls cURL and parses the result for an error
    pub fn _fetch<'a>(&self, func: &str, mut a: Easy) -> impl Future<Item = String, Error = Error> + 'a {
        let result = Arc::new(Mutex::new(Vec::new()));

        a.url(&format!("https://api.telegram.org/bot{}/{}", self.key, func)).unwrap();

        let r2 = result.clone();
        a.write_function(move |data| {
                r2.lock().unwrap().extend_from_slice(data);
                Ok(data.len())
            })
            .unwrap();

        // print debug information
        a.debug_function(|info, data| {
                match info {
                    InfoType::DataOut => {
                        println!("DataOut");
                    }
                    InfoType::Text => {
                        println!("Text");
                    }
                    InfoType::HeaderOut => {
                        println!("HeaderOut");
                    }
                    InfoType::SslDataOut => {
                        println!("SslDataOut");
                    }
                    _ => println!("something else"),
                }

                println!("{:?}", String::from_utf8_lossy(data));
            })
            .unwrap();
        //a.verbose(true).unwrap();
        //a.show_header(true).unwrap();

        self.session
            .perform(a)
            .map_err(|_| Error::TokioCurl)
            .map(move |_| {
                let response = result.lock().unwrap();
                String::from(str::from_utf8(&response).unwrap())
            })
            .and_then(|x| {
                // try to parse the result as a JSON and find the OK field.
                // If the ok field is true, then the string in "result" will be returned
                if let Ok(req) = serde_json::from_str::<Value>(&x) {
                    if let (Some(ok), res) = (req.find("ok").and_then(Value::as_bool), req.find("result")) {
                        if ok {
                            if let Some(result) = res {
                                let answer = serde_json::to_string(result).unwrap();

                                return Ok(answer);
                            }
                        }

                        match req.find("description").and_then(Value::as_str) {
                            Some(err) => Err(Error::Telegram(err.into())),
                            None => Err(Error::Telegram("Unknown".into())),
                        }
                    } else {
                        return Err(Error::JSON);
                    }
                } else {
                    return Err(Error::JSON);
                }
            })
    }
}

impl RcBot {
    /// Sets the update interval to an integer in milliseconds
    pub fn update_interval(self, interval: u64) -> RcBot {
        self.inner.update_interval.set(interval);

        self
    }

    /// Creates a new command and returns a stream which will yield a message when the command is send
    pub fn new_cmd(&self, cmd: &str) -> impl Stream<Item = (RcBot, objects::Message), Error = Error> {
        let (sender, receiver) = mpsc::unbounded();

        self.inner.handlers.borrow_mut().insert(cmd.into(), sender);

        receiver.map_err(|_| Error::Unknown)
    }

    /// Register a new commnd
    pub fn register<T>(&self, hnd: T)
        where T: Stream + 'static
    {
        self.inner.handle.spawn(hnd.for_each(|_| Ok(())).into_future().map(|_| ()).map_err(|_| ()));
    }

    /// The main update loop, the update function is called every update_interval milliseconds
    /// When an update is available the last_id will be updated and the message is filtered
    /// for commands
    /// The message is forwarded to the returned stream if no command was found
    pub fn get_stream<'a>(&'a self) -> impl Stream<Item = (RcBot, objects::Update), Error = Error> + 'a {
        use functions::*;

        Interval::new(Duration::from_millis(self.inner.update_interval.get()),
                      &self.inner.handle)
            .unwrap()
            .map_err(|_| Error::Unknown)
            .and_then(move |_| self.get_updates().offset(self.inner.last_id.get()).send())
            .map(|(_, x)| stream::iter(x.0.into_iter().map(|x| Ok(x)).collect::<Vec<Result<objects::Update, Error>>>()))
            .flatten()
            .and_then(move |x| {
                if self.inner.last_id.get() < x.update_id as u32 + 1 {
                    self.inner.last_id.set(x.update_id as u32 + 1);
                }

                Ok(x)
            })
            .filter_map(move |mut val| {
                let mut forward: Option<String> = None;

                if let Some(ref mut message) = val.message {
                    if let Some(text) = message.text.clone() {
                        let mut content = text.split_whitespace();
                        if let Some(cmd) = content.next() {
                            if self.inner.handlers.borrow_mut().contains_key(cmd) {
                                message.text = Some(content.collect::<Vec<&str>>().join(" "));

                                forward = Some(cmd.into());
                            }
                        }
                    }
                }

                if let Some(cmd) = forward {
                    if let Some(sender) = self.inner.handlers.borrow_mut().get_mut(&cmd) {
                        sender.send((self.clone(), val.message.unwrap())).unwrap();
                    }
                    return None;
                } else {
                    return Some((self.clone(), val));
                }
            })
    }

    /// helper function to start the event loop
    pub fn run<'a>(&'a self, core: &mut Core) -> Result<(), Error> {
        core.run(self.get_stream().for_each(|_| Ok(())).into_future())
    }
}
<commit_msg>add support for "/cmd@botname"<commit_after>//! This is the actual Bot module. For ergonomic reasons there is a RcBot which composes the real bot
//! as an underlying field. You should always use RcBot.

use objects;
use error::Error;

use std::str;
use std::io;
use std::time::Duration;
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::{RefCell, Cell};
use std::sync::{Arc, Mutex};

use curl::easy::{Easy, List, Form, InfoType};
use tokio_curl::Session;
use tokio_core::reactor::{Handle, Core, Interval};
use serde_json;
use serde_json::value::Value;
use futures::{Future, IntoFuture, Stream, stream};
use futures::sync::mpsc;
use futures::sync::mpsc::UnboundedSender;

/// A clonable, single threaded bot
///
/// The outer API gets implemented on RcBot
#[derive(Clone)]
pub struct RcBot {
    pub inner: Rc<Bot>,
}

impl RcBot {
    pub fn new<'a>(handle: Handle, key: &str) -> impl Future<Item = RcBot, Error = Error> + 'a {
        use functions::FunctionGetMe;
        let bot = RcBot { inner: Rc::new(Bot::new(handle, key)) };
        bot.get_me().send().map(|(mut bot, me)| {
            Rc::get_mut(&mut bot.inner).as_mut().unwrap().id = me.id;
            Rc::get_mut(&mut bot.inner).as_mut().unwrap().username = me.username.unwrap();
            bot
        })
    }
}

/// The main bot structure
pub struct Bot {
    pub key: String,
    pub handle: Handle,
    pub last_id: Cell<u32>,
    pub update_interval: Cell<u64>,
    pub handlers: RefCell<HashMap<String, UnboundedSender<(RcBot, objects::Message)>>>,
    pub session: Session,
    pub username: String,
    pub id: i64,
}

impl Bot {
    pub fn new(handle: Handle, key: &str) -> Bot {
        Bot {
            handle: handle.clone(),
            key: key.into(),
            last_id: Cell::new(0),
            update_interval: Cell::new(1000),
            handlers: RefCell::new(HashMap::new()),
            session: Session::new(handle.clone()),
            username: String::new(),
            id: 0,
        }
    }

    /// Creates a new request and adds a JSON message to it. The returned Future contains a the
    /// reply as a string.  This method should be used if no file is added because a JSON msg is
    /// always compacter than a formdata one.
    pub fn fetch_json<'a>(&self, func: &str, msg: &str) -> impl Future<Item = String, Error = Error> + 'a {
        println!("Send JSON: {}", msg);

        let mut header = List::new();
        header.append("Content-Type: application/json").unwrap();

        let mut a = Easy::new();
        a.http_headers(header).unwrap();
        a.post_fields_copy(msg.as_bytes()).unwrap();
        a.post(true).unwrap();

        self._fetch(func, a)
    }

    /// Creates a new request with some byte content (e.g. a file). The method properties have to be
    /// in the formdata setup and cannot be sent as JSON.
    pub fn fetch_formdata<'a, T>(&self,
                                 func: &str,
                                 msg: Value,
                                 mut file: T,
                                 kind: &str,
                                 file_name: &str)
                                 -> impl Future<Item = String, Error = Error> + 'a
        where T: io::Read
    {
        let mut content = Vec::new();

        let mut a = Easy::new();
        let mut form = Form::new();

        // try to read the byte content to a vector
        let _ = file.read_to_end(&mut content).unwrap();

        // add properties
        for (key, val) in msg.as_object().unwrap().iter() {
            form.part(key).contents(format!("{:?}", val).as_bytes()).add().unwrap();
        }

        // add the file
        form.part(kind).buffer(file_name, content).content_type("application/octet-stream").add().unwrap();

        // create a http post request
        a.post(true).unwrap();
        a.httppost(form).unwrap();

        self._fetch(func, a)
    }

    /// calls cURL and parses the result for an error
    pub fn _fetch<'a>(&self, func: &str, mut a: Easy) -> impl Future<Item = String, Error = Error> + 'a {
        let result = Arc::new(Mutex::new(Vec::new()));

        a.url(&format!("https://api.telegram.org/bot{}/{}", self.key, func)).unwrap();

        let r2 = result.clone();
        a.write_function(move |data| {
                r2.lock().unwrap().extend_from_slice(data);
                Ok(data.len())
            })
            .unwrap();

        // print debug information
        a.debug_function(|info, data| {
                match info {
                    InfoType::DataOut => {
                        println!("DataOut");
                    }
                    InfoType::Text => {
                        println!("Text");
                    }
                    InfoType::HeaderOut => {
                        println!("HeaderOut");
                    }
                    InfoType::SslDataOut => {
                        println!("SslDataOut");
                    }
                    _ => println!("something else"),
                }

                println!("{:?}", String::from_utf8_lossy(data));
            })
            .unwrap();
        //a.verbose(true).unwrap();
        //a.show_header(true).unwrap();

        self.session
            .perform(a)
            .map_err(|_| Error::TokioCurl)
            .map(move |_| {
                let response = result.lock().unwrap();
                String::from(str::from_utf8(&response).unwrap())
            })
            .and_then(|x| {
                // try to parse the result as a JSON and find the OK field.
                // If the ok field is true, then the string in "result" will be returned
                if let Ok(req) = serde_json::from_str::<Value>(&x) {
                    if let (Some(ok), res) = (req.find("ok").and_then(Value::as_bool), req.find("result")) {
                        if ok {
                            if let Some(result) = res {
                                let answer = serde_json::to_string(result).unwrap();

                                return Ok(answer);
                            }
                        }

                        match req.find("description").and_then(Value::as_str) {
                            Some(err) => Err(Error::Telegram(err.into())),
                            None => Err(Error::Telegram("Unknown".into())),
                        }
                    } else {
                        return Err(Error::JSON);
                    }
                } else {
                    return Err(Error::JSON);
                }
            })
    }
}

impl RcBot {
    /// Sets the update interval to an integer in milliseconds
    pub fn update_interval(self, interval: u64) -> RcBot {
        self.inner.update_interval.set(interval);

        self
    }

    /// Creates a new command and returns a stream which will yield a message when the command is send
    pub fn new_cmd(&self, cmd: &str) -> impl Stream<Item = (RcBot, objects::Message), Error = Error> {
        let (sender, receiver) = mpsc::unbounded();

        self.inner.handlers.borrow_mut().insert(cmd.into(), sender);

        receiver.map_err(|_| Error::Unknown)
    }

    /// Register a new commnd
    pub fn register<T>(&self, hnd: T)
        where T: Stream + 'static
    {
        self.inner.handle.spawn(hnd.for_each(|_| Ok(())).into_future().map(|_| ()).map_err(|_| ()));
    }

    /// The main update loop, the update function is called every update_interval milliseconds
    /// When an update is available the last_id will be updated and the message is filtered
    /// for commands
    /// The message is forwarded to the returned stream if no command was found
    pub fn get_stream<'a>(&'a self) -> impl Stream<Item = (RcBot, objects::Update), Error = Error> + 'a {
        use functions::*;

        Interval::new(Duration::from_millis(self.inner.update_interval.get()),
                      &self.inner.handle)
            .unwrap()
            .map_err(|_| Error::Unknown)
            .and_then(move |_| self.get_updates().offset(self.inner.last_id.get()).send())
            .map(|(_, x)| stream::iter(x.0.into_iter().map(|x| Ok(x)).collect::<Vec<Result<objects::Update, Error>>>()))
            .flatten()
            .and_then(move |x| {
                if self.inner.last_id.get() < x.update_id as u32 + 1 {
                    self.inner.last_id.set(x.update_id as u32 + 1);
                }

                Ok(x)
            })
            .filter_map(move |mut val| {
                let mut forward: Option<String> = None;

                if let Some(ref mut message) = val.message {
                    if let Some(text) = message.text.clone() {
                        let mut content = text.split_whitespace();
                        if let Some(cmd) = content.next() {
                            let s: Vec<&str> = cmd.split("@").take(2).collect();
                            if s.len() > 0 && (s.len() < 2 || s[1] == self.inner.username) && self.inner.handlers.borrow().contains_key(s[0]) {
                                message.text = Some(content.collect::<Vec<&str>>().join(" "));

                                forward = Some(s[0].into());
                            }
                        }
                    }
                }

                if let Some(cmd) = forward {
                    if let Some(sender) = self.inner.handlers.borrow_mut().get_mut(&cmd) {
                        sender.send((self.clone(), val.message.unwrap())).unwrap();
                    }
                    return None;
                } else {
                    return Some((self.clone(), val));
                }
            })
    }

    /// helper function to start the event loop
    pub fn run<'a>(&'a self, core: &mut Core) -> Result<(), Error> {
        core.run(self.get_stream().for_each(|_| Ok(())).into_future())
    }
}
<|endoftext|>