<commit_before>// Copyright 2015-2019 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.

use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;

use futures::lock::Mutex;
use futures::{future, Future, Poll, TryFutureExt};
use smallvec::SmallVec;

use proto::error::ProtoError;
use proto::op::ResponseCode;
use proto::xfer::{DnsHandle, DnsRequest, DnsResponse};

use crate::config::{ResolverConfig, ResolverOpts};
#[cfg(feature = "mdns")]
use crate::name_server;
use crate::name_server::{ConnectionHandle, ConnectionProvider, NameServer, StandardConnection};

/// A pool of NameServers
///
/// This is not expected to be used directly, see [`AsyncResolver`].
#[derive(Clone)]
pub struct NameServerPool<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> {
    // TODO: switch to FuturesMutex (Mutex will have some undesireable locking)
    datagram_conns: Arc<Mutex<Vec<NameServer<C, P>>>>, /* All NameServers must be the same type */
    stream_conns: Arc<Mutex<Vec<NameServer<C, P>>>>,   /* All NameServers must be the same type */
    #[cfg(feature = "mdns")]
    mdns_conns: NameServer<C, P>, /* All NameServers must be the same type */
    options: ResolverOpts,
    conn_provider: P,
}

impl NameServerPool<ConnectionHandle, StandardConnection> {
    pub(crate) fn from_config(config: &ResolverConfig, options: &ResolverOpts) -> Self {
        Self::from_config_with_provider(config, options, StandardConnection)
    }
}

impl<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> NameServerPool<C, P> {
    pub(crate) fn from_config_with_provider(
        config: &ResolverConfig,
        options: &ResolverOpts,
        conn_provider: P,
    ) -> NameServerPool<C, P> {
        let datagram_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_datagram())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        let stream_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_stream())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns)),
            stream_conns: Arc::new(Mutex::new(stream_conns)),
            #[cfg(feature = "mdns")]
            mdns_conns: name_server::mdns_nameserver(*options, conn_provider.clone()),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(not(feature = "mdns"))]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(feature = "mdns")]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        mdns_conns: NameServer<C, P>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            mdns_conns,
            options: *options,
            conn_provider,
        }
    }

    async fn try_send(
        opts: ResolverOpts,
        conns: Arc<Mutex<Vec<NameServer<C, P>>>>,
        request: DnsRequest,
    ) -> Result<DnsResponse, ProtoError> {
        let conns: Vec<NameServer<C, P>> = {
            // pull a lock on the shared connections, lock releases at the end of this scope
            let mut conns = conns.lock().await;

            // select the highest priority connection
            //   reorder the connections based on current view...
            //   this reorders the inner set
            conns.sort_unstable();

            // TODO: restrict this size to a maximum # of NameServers to try
            // get a stable view for trying all connections
            //   we split into chunks of the number of parallel requests to issue
            conns.clone()
        };

        let request_loop = request.clone();

        parallel_conn_loop(conns, request_loop, opts).await
    }
}

impl<C, P> DnsHandle for NameServerPool<C, P>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    type Response = Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send>>;

    fn send<R: Into<DnsRequest>>(&mut self, request: R) -> Self::Response {
        let opts = self.options;
        let request = request.into();
        let datagram_conns = Arc::clone(&self.datagram_conns);
        let stream_conns1 = Arc::clone(&self.stream_conns);
        let stream_conns2 = Arc::clone(&self.stream_conns);
        // TODO: remove this clone, return the Message in the error?
        let tcp_message1 = request.clone();
        let tcp_message2 = request.clone();

        // if it's a .local. query, then we *only* query mDNS, these should never be sent on to upstream resolvers
        #[cfg(feature = "mdns")]
        let mdns = mdns::maybe_local(&mut self.mdns_conns, request);

        // TODO: limited to only when mDNS is enabled, but this should probably always be enforced?
        #[cfg(not(feature = "mdns"))]
        let mdns = Local::NotMdns(request);

        // local queries are queried through mDNS
        if mdns.is_local() {
            return mdns.take_future();
        }

        // TODO: should we allow mDNS to be used for standard lookups as well?

        // it wasn't a local query, continue with standard lookup path
        let request = mdns.take_request();

        // First try the UDP connections
        Box::pin(
            Self::try_send(opts, datagram_conns, request)
                .and_then(move |response| {
                    // handling promotion from datagram to stream base on truncation in message
                    if ResponseCode::NoError == response.response_code() && response.truncated() {
                        // TCP connections should not truncate
                        future::Either::Left(Self::try_send(opts, stream_conns1, tcp_message1))
                    } else {
                        // Return the result from the UDP connection
                        future::Either::Right(future::ok(response))
                    }
                })
                // if UDP fails, try TCP
                .or_else(move |_| Self::try_send(opts, stream_conns2, tcp_message2)),
        )
    }
}

// TODO: we should be able to have a self-referential future here with Pin and not require cloned conns
/// An async function that will loop over all the conns with a max parallel request count of ops.num_concurrent_req
async fn parallel_conn_loop<C, P>(
    mut conns: Vec<NameServer<C, P>>,
    request: DnsRequest,
    opts: ResolverOpts,
) -> Result<DnsResponse, ProtoError>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    let mut err = ProtoError::from("No connections available");

    loop {
        let request_cont = request.clone();

        // construct the parallel requests, 2 is the default
        let mut par_conns = SmallVec::<[NameServer<C, P>; 2]>::new();
        let count = conns.len().min(opts.num_concurrent_reqs.max(1));
        for conn in conns.drain(..count) {
            par_conns.push(conn);
        }

        // construct the requests to send
        let requests = if par_conns.is_empty() {
            return Err(err);
        } else {
            par_conns
                .into_iter()
                .map(move |mut conn| conn.send(request_cont.clone()))
        };

        match future::select_ok(requests).await {
            Ok((sent, _)) => return Ok(sent),
            // consider a debug msg here
            Err(e) => {
                err = e;
                continue;
            }
        };
    }
}

#[cfg(feature = "mdns")]
mod mdns {
    use super::*;

    use proto::rr::domain::usage;
    use proto::DnsHandle;

    /// Returns true
    pub fn maybe_local<C, P>(name_server: &mut NameServer<C, P>, request: DnsRequest) -> Local
    where
        C: DnsHandle + 'static,
        P: ConnectionProvider<ConnHandle = C> + 'static,
    {
        if request
            .queries()
            .iter()
            .any(|query| usage::LOCAL.name().zone_of(query.name()))
        {
            Local::ResolveFuture(name_server.send(request))
        } else {
            Local::NotMdns(request)
        }
    }
}

pub enum Local {
    #[allow(dead_code)]
    ResolveFuture(Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>>),
    NotMdns(DnsRequest),
}

impl Local {
    fn is_local(&self) -> bool {
        if let Local::ResolveFuture(..) = *self {
            true
        } else {
            false
        }
    }

    /// Takes the future
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::NotMdns
    fn take_future(
        self,
    ) -> Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>> {
        match self {
            Local::ResolveFuture(future) => future,
            _ => panic!("non Local queries have no future, see take_message()"),
        }
    }

    /// Takes the message
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::ResolveFuture
    fn take_request(self) -> DnsRequest {
        match self {
            Local::NotMdns(request) => request,
            _ => panic!("Local queries must be polled, see take_future()"),
        }
    }
}

impl Future for Local {
    type Output = Result<DnsResponse, ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match *self {
            Local::ResolveFuture(ref mut ns) => ns.as_mut().poll(cx),
            // TODO: making this a panic for now
            Local::NotMdns(..) => panic!("Local queries that are not mDNS should not be polled"), //Local::NotMdns(message) => return Err(ResolveErrorKind::Message("not mDNS")),
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;

    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    use tokio::runtime::current_thread::Runtime;

    use proto::op::Query;
    use proto::rr::{Name, RecordType};
    use proto::xfer::{DnsHandle, DnsRequestOptions};

    use super::*;
    use crate::config::NameServerConfig;
    use crate::config::Protocol;

    #[ignore]
    // because of there is a real connection that needs a reasonable timeout
    #[test]
    fn test_failed_then_success_pool() {
        let config1 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 252)), 253),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let config2 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let mut resolver_config = ResolverConfig::new();
        resolver_config.add_name_server(config1);
        resolver_config.add_name_server(config2);

        let mut io_loop = Runtime::new().unwrap();
        let mut pool = NameServerPool::<_, StandardConnection>::from_config(
            &resolver_config,
            &ResolverOpts::default(),
        );

        let name = Name::parse("www.example.com.", None).unwrap();

        // TODO: it's not clear why there are two failures before the success
        for i in 0..2 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_err(),
                "iter: {}",
                i
            );
        }

        for i in 0..10 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_ok(),
                "iter: {}",
                i
            );
        }
    }
}
<commit_msg>remove extra whitespace<commit_after>// Copyright 2015-2019 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.

use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;

use futures::lock::Mutex;
use futures::{future, Future, Poll, TryFutureExt};
use smallvec::SmallVec;

use proto::error::ProtoError;
use proto::op::ResponseCode;
use proto::xfer::{DnsHandle, DnsRequest, DnsResponse};

use crate::config::{ResolverConfig, ResolverOpts};
#[cfg(feature = "mdns")]
use crate::name_server;
use crate::name_server::{ConnectionHandle, ConnectionProvider, NameServer, StandardConnection};

/// A pool of NameServers
///
/// This is not expected to be used directly, see [`AsyncResolver`].
#[derive(Clone)]
pub struct NameServerPool<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> {
    // TODO: switch to FuturesMutex (Mutex will have some undesireable locking)
    datagram_conns: Arc<Mutex<Vec<NameServer<C, P>>>>, /* All NameServers must be the same type */
    stream_conns: Arc<Mutex<Vec<NameServer<C, P>>>>,   /* All NameServers must be the same type */
    #[cfg(feature = "mdns")]
    mdns_conns: NameServer<C, P>, /* All NameServers must be the same type */
    options: ResolverOpts,
    conn_provider: P,
}

impl NameServerPool<ConnectionHandle, StandardConnection> {
    pub(crate) fn from_config(config: &ResolverConfig, options: &ResolverOpts) -> Self {
        Self::from_config_with_provider(config, options, StandardConnection)
    }
}

impl<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> NameServerPool<C, P> {
    pub(crate) fn from_config_with_provider(
        config: &ResolverConfig,
        options: &ResolverOpts,
        conn_provider: P,
    ) -> NameServerPool<C, P> {
        let datagram_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_datagram())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        let stream_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_stream())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns)),
            stream_conns: Arc::new(Mutex::new(stream_conns)),
            #[cfg(feature = "mdns")]
            mdns_conns: name_server::mdns_nameserver(*options, conn_provider.clone()),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(not(feature = "mdns"))]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(feature = "mdns")]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        mdns_conns: NameServer<C, P>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            mdns_conns,
            options: *options,
            conn_provider,
        }
    }

    async fn try_send(
        opts: ResolverOpts,
        conns: Arc<Mutex<Vec<NameServer<C, P>>>>,
        request: DnsRequest,
    ) -> Result<DnsResponse, ProtoError> {
        let conns: Vec<NameServer<C, P>> = {
            // pull a lock on the shared connections, lock releases at the end of this scope
            let mut conns = conns.lock().await;

            // select the highest priority connection
            //   reorder the connections based on current view...
            //   this reorders the inner set
            conns.sort_unstable();

            // TODO: restrict this size to a maximum # of NameServers to try
            // get a stable view for trying all connections
            //   we split into chunks of the number of parallel requests to issue
            conns.clone()
        };

        let request_loop = request.clone();

        parallel_conn_loop(conns, request_loop, opts).await
    }
}

impl<C, P> DnsHandle for NameServerPool<C, P>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    type Response = Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send>>;

    fn send<R: Into<DnsRequest>>(&mut self, request: R) -> Self::Response {
        let opts = self.options;
        let request = request.into();
        let datagram_conns = Arc::clone(&self.datagram_conns);
        let stream_conns1 = Arc::clone(&self.stream_conns);
        let stream_conns2 = Arc::clone(&self.stream_conns);
        // TODO: remove this clone, return the Message in the error?
        let tcp_message1 = request.clone();
        let tcp_message2 = request.clone();

        // if it's a .local. query, then we *only* query mDNS, these should never be sent on to upstream resolvers
        #[cfg(feature = "mdns")]
        let mdns = mdns::maybe_local(&mut self.mdns_conns, request);

        // TODO: limited to only when mDNS is enabled, but this should probably always be enforced?
        #[cfg(not(feature = "mdns"))]
        let mdns = Local::NotMdns(request);

        // local queries are queried through mDNS
        if mdns.is_local() {
            return mdns.take_future();
        }

        // TODO: should we allow mDNS to be used for standard lookups as well?

        // it wasn't a local query, continue with standard lookup path
        let request = mdns.take_request();

        // First try the UDP connections
        Box::pin(
            Self::try_send(opts, datagram_conns, request)
                .and_then(move |response| {
                    // handling promotion from datagram to stream base on truncation in message
                    if ResponseCode::NoError == response.response_code() && response.truncated() {
                        // TCP connections should not truncate
                        future::Either::Left(Self::try_send(opts, stream_conns1, tcp_message1))
                    } else {
                        // Return the result from the UDP connection
                        future::Either::Right(future::ok(response))
                    }
                })
                // if UDP fails, try TCP
                .or_else(move |_| Self::try_send(opts, stream_conns2, tcp_message2)),
        )
    }
}

// TODO: we should be able to have a self-referential future here with Pin and not require cloned conns
/// An async function that will loop over all the conns with a max parallel request count of ops.num_concurrent_req
async fn parallel_conn_loop<C, P>(
    mut conns: Vec<NameServer<C, P>>,
    request: DnsRequest,
    opts: ResolverOpts,
) -> Result<DnsResponse, ProtoError>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    let mut err = ProtoError::from("No connections available");

    loop {
        let request_cont = request.clone();

        // construct the parallel requests, 2 is the default
        let mut par_conns = SmallVec::<[NameServer<C, P>; 2]>::new();
        let count = conns.len().min(opts.num_concurrent_reqs.max(1));
        for conn in conns.drain(..count) {
            par_conns.push(conn);
        }

        // construct the requests to send
        let requests = if par_conns.is_empty() {
            return Err(err);
        } else {
            par_conns
                .into_iter()
                .map(move |mut conn| conn.send(request_cont.clone()))
        };

        match future::select_ok(requests).await {
            Ok((sent, _)) => return Ok(sent),
            // consider a debug msg here
            Err(e) => {
                err = e;
                continue;
            }
        };
    }
}

#[cfg(feature = "mdns")]
mod mdns {
    use super::*;

    use proto::rr::domain::usage;
    use proto::DnsHandle;

    /// Returns true
    pub fn maybe_local<C, P>(name_server: &mut NameServer<C, P>, request: DnsRequest) -> Local
    where
        C: DnsHandle + 'static,
        P: ConnectionProvider<ConnHandle = C> + 'static,
    {
        if request
            .queries()
            .iter()
            .any(|query| usage::LOCAL.name().zone_of(query.name()))
        {
            Local::ResolveFuture(name_server.send(request))
        } else {
            Local::NotMdns(request)
        }
    }
}

pub enum Local {
    #[allow(dead_code)]
    ResolveFuture(Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>>),
    NotMdns(DnsRequest),
}

impl Local {
    fn is_local(&self) -> bool {
        if let Local::ResolveFuture(..) = *self {
            true
        } else {
            false
        }
    }

    /// Takes the future
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::NotMdns
    fn take_future(
        self,
    ) -> Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>> {
        match self {
            Local::ResolveFuture(future) => future,
            _ => panic!("non Local queries have no future, see take_message()"),
        }
    }

    /// Takes the message
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::ResolveFuture
    fn take_request(self) -> DnsRequest {
        match self {
            Local::NotMdns(request) => request,
            _ => panic!("Local queries must be polled, see take_future()"),
        }
    }
}

impl Future for Local {
    type Output = Result<DnsResponse, ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match *self {
            Local::ResolveFuture(ref mut ns) => ns.as_mut().poll(cx),
            // TODO: making this a panic for now
            Local::NotMdns(..) => panic!("Local queries that are not mDNS should not be polled"), //Local::NotMdns(message) => return Err(ResolveErrorKind::Message("not mDNS")),
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;

    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    use tokio::runtime::current_thread::Runtime;

    use proto::op::Query;
    use proto::rr::{Name, RecordType};
    use proto::xfer::{DnsHandle, DnsRequestOptions};

    use super::*;
    use crate::config::NameServerConfig;
    use crate::config::Protocol;

    #[ignore]
    // because of there is a real connection that needs a reasonable timeout
    #[test]
    fn test_failed_then_success_pool() {
        let config1 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 252)), 253),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let config2 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let mut resolver_config = ResolverConfig::new();
        resolver_config.add_name_server(config1);
        resolver_config.add_name_server(config2);

        let mut io_loop = Runtime::new().unwrap();
        let mut pool = NameServerPool::<_, StandardConnection>::from_config(
            &resolver_config,
            &ResolverOpts::default(),
        );

        let name = Name::parse("www.example.com.", None).unwrap();

        // TODO: it's not clear why there are two failures before the success
        for i in 0..2 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_err(),
                "iter: {}",
                i
            );
        }

        for i in 0..10 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_ok(),
                "iter: {}",
                i
            );
        }
    }
}
<|endoftext|>


<commit_before>//! cargo_check provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.
use cargo_metadata::Message;
use crossbeam_channel::{never, select, unbounded, Receiver, RecvError, Sender};
use lsp_types::{
    Diagnostic, Url, WorkDoneProgress, WorkDoneProgressBegin, WorkDoneProgressEnd,
    WorkDoneProgressReport,
};
use parking_lot::RwLock;
use std::{
    collections::HashMap,
    path::PathBuf,
    process::{Command, Stdio},
    sync::Arc,
    thread::JoinHandle,
    time::Instant,
};

mod conv;

use crate::conv::{map_rust_diagnostic_to_lsp, MappedRustDiagnostic, SuggestedFix};

pub use crate::conv::url_from_path_with_drive_lowercasing;

#[derive(Clone, Debug)]
pub struct CheckOptions {
    pub enable: bool,
    pub args: Vec<String>,
    pub command: String,
    pub all_targets: bool,
}

/// CheckWatcher wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct CheckWatcher {
    pub task_recv: Receiver<CheckTask>,
    pub state: Arc<RwLock<CheckState>>,
    cmd_send: Option<Sender<CheckCommand>>,
    handle: Option<JoinHandle<()>>,
}

impl CheckWatcher {
    pub fn new(options: &CheckOptions, workspace_root: PathBuf) -> CheckWatcher {
        let options = options.clone();
        let state = Arc::new(RwLock::new(CheckState::new()));

        let (task_send, task_recv) = unbounded::<CheckTask>();
        let (cmd_send, cmd_recv) = unbounded::<CheckCommand>();
        let handle = std::thread::spawn(move || {
            let mut check = CheckWatcherState::new(options, workspace_root);
            check.run(&task_send, &cmd_recv);
        });
        CheckWatcher { task_recv, cmd_send: Some(cmd_send), handle: Some(handle), state }
    }

    /// Returns a CheckWatcher that doesn't actually do anything
    pub fn dummy() -> CheckWatcher {
        let state = Arc::new(RwLock::new(CheckState::new()));
        CheckWatcher { task_recv: never(), cmd_send: None, handle: None, state }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn update(&self) {
        if let Some(cmd_send) = &self.cmd_send {
            cmd_send.send(CheckCommand::Update).unwrap();
        }
    }
}

impl std::ops::Drop for CheckWatcher {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Take the sender out of the option
            let recv = self.cmd_send.take();

            // Dropping the sender finishes the thread loop
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

#[derive(Debug)]
pub struct CheckState {
    diagnostic_collection: HashMap<Url, Vec<Diagnostic>>,
    suggested_fix_collection: HashMap<Url, Vec<SuggestedFix>>,
}

impl CheckState {
    fn new() -> CheckState {
        CheckState {
            diagnostic_collection: HashMap::new(),
            suggested_fix_collection: HashMap::new(),
        }
    }

    /// Clear the cached diagnostics, and schedule updating diagnostics by the
    /// server, to clear stale results.
    pub fn clear(&mut self) -> Vec<Url> {
        let cleared_files: Vec<Url> = self.diagnostic_collection.keys().cloned().collect();
        self.diagnostic_collection.clear();
        self.suggested_fix_collection.clear();
        cleared_files
    }

    pub fn diagnostics_for(&self, uri: &Url) -> Option<&[Diagnostic]> {
        self.diagnostic_collection.get(uri).map(|d| d.as_slice())
    }

    pub fn fixes_for(&self, uri: &Url) -> Option<&[SuggestedFix]> {
        self.suggested_fix_collection.get(uri).map(|d| d.as_slice())
    }

    pub fn add_diagnostic_with_fixes(&mut self, file_uri: Url, diagnostic: DiagnosticWithFixes) {
        for fix in diagnostic.suggested_fixes {
            self.add_suggested_fix_for_diagnostic(fix, &diagnostic.diagnostic);
        }
        self.add_diagnostic(file_uri, diagnostic.diagnostic);
    }

    fn add_diagnostic(&mut self, file_uri: Url, diagnostic: Diagnostic) {
        let diagnostics = self.diagnostic_collection.entry(file_uri).or_default();

        // If we're building multiple targets it's possible we've already seen this diagnostic
        let is_duplicate = diagnostics.iter().any(|d| are_diagnostics_equal(d, &diagnostic));
        if is_duplicate {
            return;
        }

        diagnostics.push(diagnostic);
    }

    fn add_suggested_fix_for_diagnostic(
        &mut self,
        mut suggested_fix: SuggestedFix,
        diagnostic: &Diagnostic,
    ) {
        let file_uri = suggested_fix.location.uri.clone();
        let file_suggestions = self.suggested_fix_collection.entry(file_uri).or_default();

        let existing_suggestion: Option<&mut SuggestedFix> =
            file_suggestions.iter_mut().find(|s| s == &&suggested_fix);
        if let Some(existing_suggestion) = existing_suggestion {
            // The existing suggestion also applies to this new diagnostic
            existing_suggestion.diagnostics.push(diagnostic.clone());
        } else {
            // We haven't seen this suggestion before
            suggested_fix.diagnostics.push(diagnostic.clone());
            file_suggestions.push(suggested_fix);
        }
    }
}

#[derive(Debug)]
pub enum CheckTask {
    /// Request a clearing of all cached diagnostics from the check watcher
    ClearDiagnostics,

    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic(Url, DiagnosticWithFixes),

    /// Request check progress notification to client
    Status(WorkDoneProgress),
}

pub enum CheckCommand {
    /// Request re-start of check thread
    Update,
}

struct CheckWatcherState {
    options: CheckOptions,
    workspace_root: PathBuf,
    watcher: WatchThread,
    last_update_req: Option<Instant>,
}

impl CheckWatcherState {
    fn new(options: CheckOptions, workspace_root: PathBuf) -> CheckWatcherState {
        CheckWatcherState {
            options,
            workspace_root,
            watcher: WatchThread::dummy(),
            last_update_req: None,
        }
    }

    fn run(&mut self, task_send: &Sender<CheckTask>, cmd_recv: &Receiver<CheckCommand>) {
        loop {
            select! {
                recv(&cmd_recv) -> cmd => match cmd {
                    Ok(cmd) => self.handle_command(cmd),
                    Err(RecvError) => {
                        // Command channel has closed, so shut down
                        break;
                    },
                },
                recv(self.watcher.message_recv) -> msg => match msg {
                    Ok(msg) => self.handle_message(msg, task_send),
                    Err(RecvError) => {
                        // Watcher finished, replace it with a never channel to
                        // avoid busy-waiting.
                        std::mem::replace(&mut self.watcher.message_recv, never());
                    },
                }
            };

            if self.should_recheck() {
                self.last_update_req.take();
                task_send.send(CheckTask::ClearDiagnostics).unwrap();

                // By replacing the watcher, we drop the previous one which
                // causes it to shut down automatically.
                self.watcher = WatchThread::new(&self.options, &self.workspace_root);
            }
        }
    }

    fn should_recheck(&mut self) -> bool {
        if let Some(_last_update_req) = &self.last_update_req {
            // We currently only request an update on save, as we need up to
            // date source on disk for cargo check to do it's magic, so we
            // don't really need to debounce the requests at this point.
            return true;
        }
        false
    }

    fn handle_command(&mut self, cmd: CheckCommand) {
        match cmd {
            CheckCommand::Update => self.last_update_req = Some(Instant::now()),
        }
    }

    fn handle_message(&self, msg: CheckEvent, task_send: &Sender<CheckTask>) {
        match msg {
            CheckEvent::Begin => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Begin(WorkDoneProgressBegin {
                        title: "Running 'cargo check'".to_string(),
                        cancellable: Some(false),
                        message: None,
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::End => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd {
                        message: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerArtifact(msg)) => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Report(WorkDoneProgressReport {
                        cancellable: Some(false),
                        message: Some(msg.target.name),
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerMessage(msg)) => {
                let map_result =
                    match map_rust_diagnostic_to_lsp(&msg.message, &self.workspace_root) {
                        Some(map_result) => map_result,
                        None => return,
                    };

                let MappedRustDiagnostic { location, diagnostic, suggested_fixes } = map_result;

                let diagnostic = DiagnosticWithFixes { diagnostic, suggested_fixes };
                task_send.send(CheckTask::AddDiagnostic(location.uri, diagnostic)).unwrap();
            }

            CheckEvent::Msg(Message::BuildScriptExecuted(_msg)) => {}
            CheckEvent::Msg(Message::Unknown) => {}
        }
    }
}

#[derive(Debug)]
pub struct DiagnosticWithFixes {
    diagnostic: Diagnostic,
    suggested_fixes: Vec<SuggestedFix>,
}

/// WatchThread exists to wrap around the communication needed to be able to
/// run `cargo check` without blocking. Currently the Rust standard library
/// doesn't provide a way to read sub-process output without blocking, so we
/// have to wrap sub-processes output handling in a thread and pass messages
/// back over a channel.
/// The correct way to dispose of the thread is to drop it, on which the
/// sub-process will be killed, and the thread will be joined.
struct WatchThread {
    handle: Option<JoinHandle<()>>,
    message_recv: Receiver<CheckEvent>,
}

enum CheckEvent {
    Begin,
    Msg(cargo_metadata::Message),
    End,
}

impl WatchThread {
    fn dummy() -> WatchThread {
        WatchThread { handle: None, message_recv: never() }
    }

    fn new(options: &CheckOptions, workspace_root: &PathBuf) -> WatchThread {
        let mut args: Vec<String> = vec![
            options.command.clone(),
            "--message-format=json".to_string(),
            "--manifest-path".to_string(),
            format!("{}/Cargo.toml", workspace_root.to_string_lossy()),
        ];
        if options.all_targets {
            args.push("--all-targets".to_string());
        }
        args.extend(options.args.iter().cloned());

        let (message_send, message_recv) = unbounded();
        let enabled = options.enable;
        let handle = std::thread::spawn(move || {
            if !enabled {
                return;
            }

            let mut command = Command::new("cargo")
                .args(&args)
                .stdout(Stdio::piped())
                .stderr(Stdio::null())
                .spawn()
                .expect("couldn't launch cargo");

            // If we trigger an error here, we will do so in the loop instead,
            // which will break out of the loop, and continue the shutdown
            let _ = message_send.send(CheckEvent::Begin);

            for message in cargo_metadata::parse_messages(command.stdout.take().unwrap()) {
                let message = match message {
                    Ok(message) => message,
                    Err(err) => {
                        log::error!("Invalid json from cargo check, ignoring: {}", err);
                        continue;
                    }
                };

                match message_send.send(CheckEvent::Msg(message)) {
                    Ok(()) => {}
                    Err(_err) => {
                        // The send channel was closed, so we want to shutdown
                        break;
                    }
                }
            }

            // We can ignore any error here, as we are already in the progress
            // of shutting down.
            let _ = message_send.send(CheckEvent::End);

            // It is okay to ignore the result, as it only errors if the process is already dead
            let _ = command.kill();

            // Again, we don't care about the exit status so just ignore the result
            let _ = command.wait();
        });
        WatchThread { handle: Some(handle), message_recv }
    }
}

impl std::ops::Drop for WatchThread {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Replace our reciever with dummy one, so we can drop and close the
            // one actually communicating with the thread
            let recv = std::mem::replace(&mut self.message_recv, never());

            // Dropping the original reciever initiates thread sub-process shutdown
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

fn are_diagnostics_equal(left: &Diagnostic, right: &Diagnostic) -> bool {
    left.source == right.source
        && left.severity == right.severity
        && left.range == right.range
        && left.message == right.message
}
<commit_msg>Tweak naming slightly<commit_after>//! cargo_check provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.
use cargo_metadata::Message;
use crossbeam_channel::{never, select, unbounded, Receiver, RecvError, Sender};
use lsp_types::{
    Diagnostic, Url, WorkDoneProgress, WorkDoneProgressBegin, WorkDoneProgressEnd,
    WorkDoneProgressReport,
};
use parking_lot::RwLock;
use std::{
    collections::HashMap,
    path::PathBuf,
    process::{Command, Stdio},
    sync::Arc,
    thread::JoinHandle,
    time::Instant,
};

mod conv;

use crate::conv::{map_rust_diagnostic_to_lsp, MappedRustDiagnostic, SuggestedFix};

pub use crate::conv::url_from_path_with_drive_lowercasing;

#[derive(Clone, Debug)]
pub struct CheckOptions {
    pub enable: bool,
    pub args: Vec<String>,
    pub command: String,
    pub all_targets: bool,
}

/// CheckWatcher wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct CheckWatcher {
    pub task_recv: Receiver<CheckTask>,
    pub state: Arc<RwLock<CheckState>>,
    cmd_send: Option<Sender<CheckCommand>>,
    handle: Option<JoinHandle<()>>,
}

impl CheckWatcher {
    pub fn new(options: &CheckOptions, workspace_root: PathBuf) -> CheckWatcher {
        let options = options.clone();
        let state = Arc::new(RwLock::new(CheckState::new()));

        let (task_send, task_recv) = unbounded::<CheckTask>();
        let (cmd_send, cmd_recv) = unbounded::<CheckCommand>();
        let handle = std::thread::spawn(move || {
            let mut check = CheckWatcherThread::new(options, workspace_root);
            check.run(&task_send, &cmd_recv);
        });
        CheckWatcher { task_recv, cmd_send: Some(cmd_send), handle: Some(handle), state }
    }

    /// Returns a CheckWatcher that doesn't actually do anything
    pub fn dummy() -> CheckWatcher {
        let state = Arc::new(RwLock::new(CheckState::new()));
        CheckWatcher { task_recv: never(), cmd_send: None, handle: None, state }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn update(&self) {
        if let Some(cmd_send) = &self.cmd_send {
            cmd_send.send(CheckCommand::Update).unwrap();
        }
    }
}

impl std::ops::Drop for CheckWatcher {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Take the sender out of the option
            let recv = self.cmd_send.take();

            // Dropping the sender finishes the thread loop
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

#[derive(Debug)]
pub struct CheckState {
    diagnostic_collection: HashMap<Url, Vec<Diagnostic>>,
    suggested_fix_collection: HashMap<Url, Vec<SuggestedFix>>,
}

impl CheckState {
    fn new() -> CheckState {
        CheckState {
            diagnostic_collection: HashMap::new(),
            suggested_fix_collection: HashMap::new(),
        }
    }

    /// Clear the cached diagnostics, and schedule updating diagnostics by the
    /// server, to clear stale results.
    pub fn clear(&mut self) -> Vec<Url> {
        let cleared_files: Vec<Url> = self.diagnostic_collection.keys().cloned().collect();
        self.diagnostic_collection.clear();
        self.suggested_fix_collection.clear();
        cleared_files
    }

    pub fn diagnostics_for(&self, uri: &Url) -> Option<&[Diagnostic]> {
        self.diagnostic_collection.get(uri).map(|d| d.as_slice())
    }

    pub fn fixes_for(&self, uri: &Url) -> Option<&[SuggestedFix]> {
        self.suggested_fix_collection.get(uri).map(|d| d.as_slice())
    }

    pub fn add_diagnostic_with_fixes(&mut self, file_uri: Url, diagnostic: DiagnosticWithFixes) {
        for fix in diagnostic.suggested_fixes {
            self.add_suggested_fix_for_diagnostic(fix, &diagnostic.diagnostic);
        }
        self.add_diagnostic(file_uri, diagnostic.diagnostic);
    }

    fn add_diagnostic(&mut self, file_uri: Url, diagnostic: Diagnostic) {
        let diagnostics = self.diagnostic_collection.entry(file_uri).or_default();

        // If we're building multiple targets it's possible we've already seen this diagnostic
        let is_duplicate = diagnostics.iter().any(|d| are_diagnostics_equal(d, &diagnostic));
        if is_duplicate {
            return;
        }

        diagnostics.push(diagnostic);
    }

    fn add_suggested_fix_for_diagnostic(
        &mut self,
        mut suggested_fix: SuggestedFix,
        diagnostic: &Diagnostic,
    ) {
        let file_uri = suggested_fix.location.uri.clone();
        let file_suggestions = self.suggested_fix_collection.entry(file_uri).or_default();

        let existing_suggestion: Option<&mut SuggestedFix> =
            file_suggestions.iter_mut().find(|s| s == &&suggested_fix);
        if let Some(existing_suggestion) = existing_suggestion {
            // The existing suggestion also applies to this new diagnostic
            existing_suggestion.diagnostics.push(diagnostic.clone());
        } else {
            // We haven't seen this suggestion before
            suggested_fix.diagnostics.push(diagnostic.clone());
            file_suggestions.push(suggested_fix);
        }
    }
}

#[derive(Debug)]
pub enum CheckTask {
    /// Request a clearing of all cached diagnostics from the check watcher
    ClearDiagnostics,

    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic(Url, DiagnosticWithFixes),

    /// Request check progress notification to client
    Status(WorkDoneProgress),
}

pub enum CheckCommand {
    /// Request re-start of check thread
    Update,
}

struct CheckWatcherThread {
    options: CheckOptions,
    workspace_root: PathBuf,
    watcher: WatchThread,
    last_update_req: Option<Instant>,
}

impl CheckWatcherThread {
    fn new(options: CheckOptions, workspace_root: PathBuf) -> CheckWatcherThread {
        CheckWatcherThread {
            options,
            workspace_root,
            watcher: WatchThread::dummy(),
            last_update_req: None,
        }
    }

    fn run(&mut self, task_send: &Sender<CheckTask>, cmd_recv: &Receiver<CheckCommand>) {
        loop {
            select! {
                recv(&cmd_recv) -> cmd => match cmd {
                    Ok(cmd) => self.handle_command(cmd),
                    Err(RecvError) => {
                        // Command channel has closed, so shut down
                        break;
                    },
                },
                recv(self.watcher.message_recv) -> msg => match msg {
                    Ok(msg) => self.handle_message(msg, task_send),
                    Err(RecvError) => {
                        // Watcher finished, replace it with a never channel to
                        // avoid busy-waiting.
                        std::mem::replace(&mut self.watcher.message_recv, never());
                    },
                }
            };

            if self.should_recheck() {
                self.last_update_req.take();
                task_send.send(CheckTask::ClearDiagnostics).unwrap();

                // By replacing the watcher, we drop the previous one which
                // causes it to shut down automatically.
                self.watcher = WatchThread::new(&self.options, &self.workspace_root);
            }
        }
    }

    fn should_recheck(&mut self) -> bool {
        if let Some(_last_update_req) = &self.last_update_req {
            // We currently only request an update on save, as we need up to
            // date source on disk for cargo check to do it's magic, so we
            // don't really need to debounce the requests at this point.
            return true;
        }
        false
    }

    fn handle_command(&mut self, cmd: CheckCommand) {
        match cmd {
            CheckCommand::Update => self.last_update_req = Some(Instant::now()),
        }
    }

    fn handle_message(&self, msg: CheckEvent, task_send: &Sender<CheckTask>) {
        match msg {
            CheckEvent::Begin => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Begin(WorkDoneProgressBegin {
                        title: "Running 'cargo check'".to_string(),
                        cancellable: Some(false),
                        message: None,
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::End => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::End(WorkDoneProgressEnd {
                        message: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerArtifact(msg)) => {
                task_send
                    .send(CheckTask::Status(WorkDoneProgress::Report(WorkDoneProgressReport {
                        cancellable: Some(false),
                        message: Some(msg.target.name),
                        percentage: None,
                    })))
                    .unwrap();
            }

            CheckEvent::Msg(Message::CompilerMessage(msg)) => {
                let map_result =
                    match map_rust_diagnostic_to_lsp(&msg.message, &self.workspace_root) {
                        Some(map_result) => map_result,
                        None => return,
                    };

                let MappedRustDiagnostic { location, diagnostic, suggested_fixes } = map_result;

                let diagnostic = DiagnosticWithFixes { diagnostic, suggested_fixes };
                task_send.send(CheckTask::AddDiagnostic(location.uri, diagnostic)).unwrap();
            }

            CheckEvent::Msg(Message::BuildScriptExecuted(_msg)) => {}
            CheckEvent::Msg(Message::Unknown) => {}
        }
    }
}

#[derive(Debug)]
pub struct DiagnosticWithFixes {
    diagnostic: Diagnostic,
    suggested_fixes: Vec<SuggestedFix>,
}

/// WatchThread exists to wrap around the communication needed to be able to
/// run `cargo check` without blocking. Currently the Rust standard library
/// doesn't provide a way to read sub-process output without blocking, so we
/// have to wrap sub-processes output handling in a thread and pass messages
/// back over a channel.
/// The correct way to dispose of the thread is to drop it, on which the
/// sub-process will be killed, and the thread will be joined.
struct WatchThread {
    handle: Option<JoinHandle<()>>,
    message_recv: Receiver<CheckEvent>,
}

enum CheckEvent {
    Begin,
    Msg(cargo_metadata::Message),
    End,
}

impl WatchThread {
    fn dummy() -> WatchThread {
        WatchThread { handle: None, message_recv: never() }
    }

    fn new(options: &CheckOptions, workspace_root: &PathBuf) -> WatchThread {
        let mut args: Vec<String> = vec![
            options.command.clone(),
            "--message-format=json".to_string(),
            "--manifest-path".to_string(),
            format!("{}/Cargo.toml", workspace_root.to_string_lossy()),
        ];
        if options.all_targets {
            args.push("--all-targets".to_string());
        }
        args.extend(options.args.iter().cloned());

        let (message_send, message_recv) = unbounded();
        let enabled = options.enable;
        let handle = std::thread::spawn(move || {
            if !enabled {
                return;
            }

            let mut command = Command::new("cargo")
                .args(&args)
                .stdout(Stdio::piped())
                .stderr(Stdio::null())
                .spawn()
                .expect("couldn't launch cargo");

            // If we trigger an error here, we will do so in the loop instead,
            // which will break out of the loop, and continue the shutdown
            let _ = message_send.send(CheckEvent::Begin);

            for message in cargo_metadata::parse_messages(command.stdout.take().unwrap()) {
                let message = match message {
                    Ok(message) => message,
                    Err(err) => {
                        log::error!("Invalid json from cargo check, ignoring: {}", err);
                        continue;
                    }
                };

                match message_send.send(CheckEvent::Msg(message)) {
                    Ok(()) => {}
                    Err(_err) => {
                        // The send channel was closed, so we want to shutdown
                        break;
                    }
                }
            }

            // We can ignore any error here, as we are already in the progress
            // of shutting down.
            let _ = message_send.send(CheckEvent::End);

            // It is okay to ignore the result, as it only errors if the process is already dead
            let _ = command.kill();

            // Again, we don't care about the exit status so just ignore the result
            let _ = command.wait();
        });
        WatchThread { handle: Some(handle), message_recv }
    }
}

impl std::ops::Drop for WatchThread {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            // Replace our reciever with dummy one, so we can drop and close the
            // one actually communicating with the thread
            let recv = std::mem::replace(&mut self.message_recv, never());

            // Dropping the original reciever initiates thread sub-process shutdown
            drop(recv);

            // Join the thread, it should finish shortly. We don't really care
            // whether it panicked, so it is safe to ignore the result
            let _ = handle.join();
        }
    }
}

fn are_diagnostics_equal(left: &Diagnostic, right: &Diagnostic) -> bool {
    left.source == right.source
        && left.severity == right.severity
        && left.range == right.range
        && left.message == right.message
}
<|endoftext|>


<commit_before>use std::ptr;

use log::{debug, info, warn};
use failure::format_err;
use openldap::{self, LDAPResponse, RustLDAP};
use regex::Regex;

use octobot_lib::config::LdapConfig;
use octobot_lib::errors::*;

pub struct LDAPEntry {
    pub dn: String,
}

fn new_ldap(url: &str) -> Result<RustLDAP> {
    let ldap = RustLDAP::new(url)?;

    ldap.set_option(
        openldap::codes::options::LDAP_OPT_PROTOCOL_VERSION,
        &openldap::codes::versions::LDAP_VERSION3,
    );

    ldap.set_option(
        openldap::codes::options::LDAP_OPT_X_TLS_REQUIRE_CERT,
        &openldap::codes::options::LDAP_OPT_X_TLS_DEMAND,
    );

    Ok(ldap)
}

pub fn auth(user: &str, pass: &str, config: &LdapConfig) -> Result<bool> {
    if user.is_empty() {
        info!("Cannot authenticate without username");
        return Ok(false);
    }

    // in the absence of `ldap_escape` from ldap3, just whitelist acceptable characters
    let re = Regex::new(r"([^A-Za-z0-9\.\-_@])").unwrap();
    for cap in re.captures_iter(user) {
        info!("Invalid username character in username: '{}', '{}'", &cap[1], user);
        return Ok(false);
    }

    let user_filters = config.userid_attributes.iter().map(|a| format!("({}={})", a, user)).collect::<Vec<_>>();

    let user_filter;
    if user_filters.len() == 0 {
        info!("Cannot authenticate without userid attributes");
        return Ok(false);
    } else if user_filters.len() == 1 {
        user_filter = user_filters[0].clone();
    } else {
        user_filter = format!("(|{})", user_filters.join(""));
    }

    // search for the user's DN
    let results = search(config, Some(&user_filter), 1)?;

    if results.is_empty() {
        debug!("No users found matching {}", user);
        return Ok(false);
    }
    if results.len() > 1 {
        info!("Too many users found matching {}", user);
        return Ok(false);
    }

    let user_dn = &results[0].dn;
    if user_dn.is_empty() {
        info!("User found but with empty DN!");
        return Ok(false);
    }

    // now try to bind as the user
    let ldap = new_ldap(&config.url)?;
    let res = ldap.simple_bind(&user_dn, &pass)?;
    if res == 0 {
        Ok(true)
    } else if res == 49 {
        // Avoid error messages for invalid creds
        Ok(false)
    } else {
        info!("LDAP auth failed with error code {}", res);
        Ok(false)
    }
}

pub fn search(config: &LdapConfig, extra_filter: Option<&str>, max_results: i32) -> Result<Vec<LDAPEntry>> {
    let ldap = new_ldap(&config.url)?;
    let bind_res = ldap.simple_bind(&config.bind_user, &config.bind_pass)?;
    if bind_res != 0 {
        return Err(format_err!("LDAP service account bind failed with error code {}", bind_res));
    }

    let mut search_filters: Vec<String> = vec![];
    if let Some(f) = extra_filter {
        search_filters.push(f.to_string());
    }
    if let Some(ref f) = config.search_filter {
        search_filters.push(f.to_string());
    }

    let search_filter;
    if search_filters.is_empty() {
        warn!("No LDAP search filter configured. There may be lots of results");
        search_filter = "(objectClass=*)".to_string();
    } else if search_filters.len() == 1 {
        search_filter = search_filters[0].clone();
    } else {
        search_filter = format!("(&{})", search_filters.join(""));
    }

    let resp: Result<LDAPResponse> = ldap.ldap_search(
        &config.base_dn,
        openldap::codes::scopes::LDAP_SCOPE_SUB,
        Some(&search_filter),
        None, // attrs
        false, // attrsonly
        None, // server controls
        None, // client controls
        ptr::null_mut(), // timeout
        max_results,
    ).map_err(|e| format_err!("Error on LDAP search: {}", e));

    let entries = resp?.into_iter()
        .filter_map(|attrs| {
            let dn = attrs.get("dn").unwrap_or(&vec![]).iter().next().map(|s| s.to_string()).unwrap_or(
                String::new(),
            );
            if dn.is_empty() {
                warn!("Found entry with empty DN! Skipping.");
                None
            } else {
                Some(LDAPEntry { dn: dn })
            }
        })
        .collect::<Vec<LDAPEntry>>();

    Ok(entries)
}
<commit_msg>Require password for login (#288)<commit_after>use std::ptr;

use log::{debug, info, warn};
use failure::format_err;
use openldap::{self, LDAPResponse, RustLDAP};
use regex::Regex;

use octobot_lib::config::LdapConfig;
use octobot_lib::errors::*;

pub struct LDAPEntry {
    pub dn: String,
}

fn new_ldap(url: &str) -> Result<RustLDAP> {
    let ldap = RustLDAP::new(url)?;

    ldap.set_option(
        openldap::codes::options::LDAP_OPT_PROTOCOL_VERSION,
        &openldap::codes::versions::LDAP_VERSION3,
    );

    ldap.set_option(
        openldap::codes::options::LDAP_OPT_X_TLS_REQUIRE_CERT,
        &openldap::codes::options::LDAP_OPT_X_TLS_DEMAND,
    );

    Ok(ldap)
}

pub fn auth(user: &str, pass: &str, config: &LdapConfig) -> Result<bool> {
    if user.is_empty() {
        info!("Cannot authenticate without username");
        return Ok(false);
    }
    if pass.is_empty() {
        info!("Cannot authenticate without password");
        return Ok(false);
    }

    // in the absence of `ldap_escape` from ldap3, just whitelist acceptable characters
    let re = Regex::new(r"([^A-Za-z0-9\.\-_@])").unwrap();
    for cap in re.captures_iter(user) {
        info!("Invalid username character in username: '{}', '{}'", &cap[1], user);
        return Ok(false);
    }

    let user_filters = config.userid_attributes.iter().map(|a| format!("({}={})", a, user)).collect::<Vec<_>>();

    let user_filter;
    if user_filters.len() == 0 {
        info!("Cannot authenticate without userid attributes");
        return Ok(false);
    } else if user_filters.len() == 1 {
        user_filter = user_filters[0].clone();
    } else {
        user_filter = format!("(|{})", user_filters.join(""));
    }

    // search for the user's DN
    let results = search(config, Some(&user_filter), 1)?;

    if results.is_empty() {
        debug!("No users found matching {}", user);
        return Ok(false);
    }
    if results.len() > 1 {
        info!("Too many users found matching {}", user);
        return Ok(false);
    }

    let user_dn = &results[0].dn;
    if user_dn.is_empty() {
        info!("User found but with empty DN!");
        return Ok(false);
    }

    // now try to bind as the user
    let ldap = new_ldap(&config.url)?;
    let res = ldap.simple_bind(&user_dn, &pass)?;
    if res == 0 {
        Ok(true)
    } else if res == 49 {
        // Avoid error messages for invalid creds
        Ok(false)
    } else {
        info!("LDAP auth failed with error code {}", res);
        Ok(false)
    }
}

pub fn search(config: &LdapConfig, extra_filter: Option<&str>, max_results: i32) -> Result<Vec<LDAPEntry>> {
    let ldap = new_ldap(&config.url)?;
    let bind_res = ldap.simple_bind(&config.bind_user, &config.bind_pass)?;
    if bind_res != 0 {
        return Err(format_err!("LDAP service account bind failed with error code {}", bind_res));
    }

    let mut search_filters: Vec<String> = vec![];
    if let Some(f) = extra_filter {
        search_filters.push(f.to_string());
    }
    if let Some(ref f) = config.search_filter {
        search_filters.push(f.to_string());
    }

    let search_filter;
    if search_filters.is_empty() {
        warn!("No LDAP search filter configured. There may be lots of results");
        search_filter = "(objectClass=*)".to_string();
    } else if search_filters.len() == 1 {
        search_filter = search_filters[0].clone();
    } else {
        search_filter = format!("(&{})", search_filters.join(""));
    }

    let resp: Result<LDAPResponse> = ldap.ldap_search(
        &config.base_dn,
        openldap::codes::scopes::LDAP_SCOPE_SUB,
        Some(&search_filter),
        None, // attrs
        false, // attrsonly
        None, // server controls
        None, // client controls
        ptr::null_mut(), // timeout
        max_results,
    ).map_err(|e| format_err!("Error on LDAP search: {}", e));

    let entries = resp?.into_iter()
        .filter_map(|attrs| {
            let dn = attrs.get("dn").unwrap_or(&vec![]).iter().next().map(|s| s.to_string()).unwrap_or(
                String::new(),
            );
            if dn.is_empty() {
                warn!("Found entry with empty DN! Skipping.");
                None
            } else {
                Some(LDAPEntry { dn: dn })
            }
        })
        .collect::<Vec<LDAPEntry>>();

    Ok(entries)
}
<|endoftext|>


<commit_before>//! checks for attributes

use reexport::*;
use rustc::lint::*;
use rustc::hir::*;
use semver::Version;
use syntax::ast::{Attribute, Lit, LitKind, MetaItemKind};
use syntax::codemap::Span;
use utils::{in_macro, match_path, span_lint, span_lint_and_then, snippet_opt};
use utils::paths;

/// **What it does:** Checks for items annotated with `#[inline(always)]`,
/// unless the annotated function is empty or simply panics.
///
/// **Why is this bad?** While there are valid uses of this annotation (and once
/// you know when to use it, by all means `allow` this lint), it's a common
/// newbie-mistake to pepper one's code with it.
///
/// As a rule of thumb, before slapping `#[inline(always)]` on a function,
/// measure if that additional function call really affects your runtime profile
/// sufficiently to make up for the increase in compile time.
///
/// **Known problems:** False positives, big time. This lint is meant to be
/// deactivated by everyone doing serious performance work. This means having
/// done the measurement.
///
/// **Example:**
/// ```rust
/// #[inline(always)]
/// fn not_quite_hot_code(..) { ... }
/// ```
declare_lint! {
    pub INLINE_ALWAYS,
    Warn,
    "use of `#[inline(always)]`"
}

/// **What it does:** Checks for `extern crate` and `use` items annotated with lint attributes
///
/// **Why is this bad?** Lint attributes have no effect on crate imports. Most likely a `!` was forgotten
///
/// **Known problems:** Technically one might allow `unused_import` on a `use` item, but it's easier to remove the unused item.
///
/// **Example:**
/// ```rust
/// #[deny(dead_code)]
/// extern crate foo;
/// #[allow(unused_import)]
/// use foo::bar;
/// ```
declare_lint! {
    pub USELESS_ATTRIBUTE,
    Warn,
    "use of lint attributes on `extern crate` items"
}

/// **What it does:** Checks for `#[deprecated]` annotations with a `since`
/// field that is not a valid semantic version.
///
/// **Why is this bad?** For checking the version of the deprecation, it must be
/// a valid semver. Failing that, the contained information is useless.
///
/// **Known problems:** None.
///
/// **Example:**
/// ```rust
/// #[deprecated(since = "forever")]
/// fn something_else(..) { ... }
/// ```
declare_lint! {
    pub DEPRECATED_SEMVER,
    Warn,
    "use of `#[deprecated(since = \"x\")]` where x is not semver"
}

#[derive(Copy,Clone)]
pub struct AttrPass;

impl LintPass for AttrPass {
    fn get_lints(&self) -> LintArray {
        lint_array!(INLINE_ALWAYS, DEPRECATED_SEMVER, USELESS_ATTRIBUTE)
    }
}

impl LateLintPass for AttrPass {
    fn check_attribute(&mut self, cx: &LateContext, attr: &Attribute) {
        if let MetaItemKind::List(ref name, ref items) = attr.node.value.node {
            if items.is_empty() || name != &"deprecated" {
                return;
            }
            for item in items {
                if let MetaItemKind::NameValue(ref name, ref lit) = item.node {
                    if name == &"since" {
                        check_semver(cx, item.span, lit);
                    }
                }
            }
        }
    }

    fn check_item(&mut self, cx: &LateContext, item: &Item) {
        if is_relevant_item(item) {
            check_attrs(cx, item.span, &item.name, &item.attrs)
        }
        match item.node {
            ItemExternCrate(_) |
            ItemUse(_) => {
                for attr in &item.attrs {
                    if let MetaItemKind::List(ref name, _) = attr.node.value.node {
                        match &**name {
                            "allow" | "warn" | "deny" | "forbid" => {
                                span_lint_and_then(cx, USELESS_ATTRIBUTE, attr.span,
                                                   "useless lint attribute",
                                                   |db| {
                                    if let Some(mut sugg) = snippet_opt(cx, attr.span) {
                                        sugg.insert(1, '!');
                                        db.span_suggestion(attr.span, "if you just forgot a `!`, use", sugg);
                                    }
                                });
                            },
                            _ => {},
                        }
                    }
                }
            },
            _ => {},
        }
    }

    fn check_impl_item(&mut self, cx: &LateContext, item: &ImplItem) {
        if is_relevant_impl(item) {
            check_attrs(cx, item.span, &item.name, &item.attrs)
        }
    }

    fn check_trait_item(&mut self, cx: &LateContext, item: &TraitItem) {
        if is_relevant_trait(item) {
            check_attrs(cx, item.span, &item.name, &item.attrs)
        }
    }
}

fn is_relevant_item(item: &Item) -> bool {
    if let ItemFn(_, _, _, _, _, ref block) = item.node {
        is_relevant_block(block)
    } else {
        false
    }
}

fn is_relevant_impl(item: &ImplItem) -> bool {
    match item.node {
        ImplItemKind::Method(_, ref block) => is_relevant_block(block),
        _ => false,
    }
}

fn is_relevant_trait(item: &TraitItem) -> bool {
    match item.node {
        MethodTraitItem(_, None) => true,
        MethodTraitItem(_, Some(ref block)) => is_relevant_block(block),
        _ => false,
    }
}

fn is_relevant_block(block: &Block) -> bool {
    for stmt in &block.stmts {
        match stmt.node {
            StmtDecl(_, _) => return true,
            StmtExpr(ref expr, _) |
            StmtSemi(ref expr, _) => {
                return is_relevant_expr(expr);
            }
        }
    }
    block.expr.as_ref().map_or(false, |e| is_relevant_expr(e))
}

fn is_relevant_expr(expr: &Expr) -> bool {
    match expr.node {
        ExprBlock(ref block) => is_relevant_block(block),
        ExprRet(Some(ref e)) => is_relevant_expr(e),
        ExprRet(None) | ExprBreak(_) => false,
        ExprCall(ref path_expr, _) => {
            if let ExprPath(_, ref path) = path_expr.node {
                !match_path(path, &paths::BEGIN_PANIC)
            } else {
                true
            }
        }
        _ => true,
    }
}

fn check_attrs(cx: &LateContext, span: Span, name: &Name, attrs: &[Attribute]) {
    if in_macro(cx, span) {
        return;
    }

    for attr in attrs {
        if let MetaItemKind::List(ref inline, ref values) = attr.node.value.node {
            if values.len() != 1 || inline != &"inline" {
                continue;
            }
            if let MetaItemKind::Word(ref always) = values[0].node {
                if always != &"always" {
                    continue;
                }
                span_lint(cx,
                          INLINE_ALWAYS,
                          attr.span,
                          &format!("you have declared `#[inline(always)]` on `{}`. This is usually a bad idea",
                                   name));
            }
        }
    }
}

fn check_semver(cx: &LateContext, span: Span, lit: &Lit) {
    if let LitKind::Str(ref is, _) = lit.node {
        if Version::parse(&*is).is_ok() {
            return;
        }
    }
    span_lint(cx,
              DEPRECATED_SEMVER,
              span,
              "the since field must contain a semver-compliant version");
}
<commit_msg>don't lint if snippet is shortert than 1 character<commit_after>//! checks for attributes

use reexport::*;
use rustc::lint::*;
use rustc::hir::*;
use semver::Version;
use syntax::ast::{Attribute, Lit, LitKind, MetaItemKind};
use syntax::codemap::Span;
use utils::{in_macro, match_path, span_lint, span_lint_and_then, snippet_opt};
use utils::paths;

/// **What it does:** Checks for items annotated with `#[inline(always)]`,
/// unless the annotated function is empty or simply panics.
///
/// **Why is this bad?** While there are valid uses of this annotation (and once
/// you know when to use it, by all means `allow` this lint), it's a common
/// newbie-mistake to pepper one's code with it.
///
/// As a rule of thumb, before slapping `#[inline(always)]` on a function,
/// measure if that additional function call really affects your runtime profile
/// sufficiently to make up for the increase in compile time.
///
/// **Known problems:** False positives, big time. This lint is meant to be
/// deactivated by everyone doing serious performance work. This means having
/// done the measurement.
///
/// **Example:**
/// ```rust
/// #[inline(always)]
/// fn not_quite_hot_code(..) { ... }
/// ```
declare_lint! {
    pub INLINE_ALWAYS,
    Warn,
    "use of `#[inline(always)]`"
}

/// **What it does:** Checks for `extern crate` and `use` items annotated with lint attributes
///
/// **Why is this bad?** Lint attributes have no effect on crate imports. Most likely a `!` was forgotten
///
/// **Known problems:** Technically one might allow `unused_import` on a `use` item, but it's easier to remove the unused item.
///
/// **Example:**
/// ```rust
/// #[deny(dead_code)]
/// extern crate foo;
/// #[allow(unused_import)]
/// use foo::bar;
/// ```
declare_lint! {
    pub USELESS_ATTRIBUTE,
    Warn,
    "use of lint attributes on `extern crate` items"
}

/// **What it does:** Checks for `#[deprecated]` annotations with a `since`
/// field that is not a valid semantic version.
///
/// **Why is this bad?** For checking the version of the deprecation, it must be
/// a valid semver. Failing that, the contained information is useless.
///
/// **Known problems:** None.
///
/// **Example:**
/// ```rust
/// #[deprecated(since = "forever")]
/// fn something_else(..) { ... }
/// ```
declare_lint! {
    pub DEPRECATED_SEMVER,
    Warn,
    "use of `#[deprecated(since = \"x\")]` where x is not semver"
}

#[derive(Copy,Clone)]
pub struct AttrPass;

impl LintPass for AttrPass {
    fn get_lints(&self) -> LintArray {
        lint_array!(INLINE_ALWAYS, DEPRECATED_SEMVER, USELESS_ATTRIBUTE)
    }
}

impl LateLintPass for AttrPass {
    fn check_attribute(&mut self, cx: &LateContext, attr: &Attribute) {
        if let MetaItemKind::List(ref name, ref items) = attr.node.value.node {
            if items.is_empty() || name != &"deprecated" {
                return;
            }
            for item in items {
                if let MetaItemKind::NameValue(ref name, ref lit) = item.node {
                    if name == &"since" {
                        check_semver(cx, item.span, lit);
                    }
                }
            }
        }
    }

    fn check_item(&mut self, cx: &LateContext, item: &Item) {
        if is_relevant_item(item) {
            check_attrs(cx, item.span, &item.name, &item.attrs)
        }
        match item.node {
            ItemExternCrate(_) |
            ItemUse(_) => {
                for attr in &item.attrs {
                    if let MetaItemKind::List(ref name, _) = attr.node.value.node {
                        match &**name {
                            "allow" | "warn" | "deny" | "forbid" => {
                                if let Some(mut sugg) = snippet_opt(cx, attr.span) {
                                    if sugg.len() > 1 {
                                        span_lint_and_then(cx, USELESS_ATTRIBUTE, attr.span,
                                                           "useless lint attribute",
                                                           |db| {
                                            sugg.insert(1, '!');
                                            db.span_suggestion(attr.span, "if you just forgot a `!`, use", sugg);
                                        });
                                    }
                                }
                            },
                            _ => {},
                        }
                    }
                }
            },
            _ => {},
        }
    }

    fn check_impl_item(&mut self, cx: &LateContext, item: &ImplItem) {
        if is_relevant_impl(item) {
            check_attrs(cx, item.span, &item.name, &item.attrs)
        }
    }

    fn check_trait_item(&mut self, cx: &LateContext, item: &TraitItem) {
        if is_relevant_trait(item) {
            check_attrs(cx, item.span, &item.name, &item.attrs)
        }
    }
}

fn is_relevant_item(item: &Item) -> bool {
    if let ItemFn(_, _, _, _, _, ref block) = item.node {
        is_relevant_block(block)
    } else {
        false
    }
}

fn is_relevant_impl(item: &ImplItem) -> bool {
    match item.node {
        ImplItemKind::Method(_, ref block) => is_relevant_block(block),
        _ => false,
    }
}

fn is_relevant_trait(item: &TraitItem) -> bool {
    match item.node {
        MethodTraitItem(_, None) => true,
        MethodTraitItem(_, Some(ref block)) => is_relevant_block(block),
        _ => false,
    }
}

fn is_relevant_block(block: &Block) -> bool {
    for stmt in &block.stmts {
        match stmt.node {
            StmtDecl(_, _) => return true,
            StmtExpr(ref expr, _) |
            StmtSemi(ref expr, _) => {
                return is_relevant_expr(expr);
            }
        }
    }
    block.expr.as_ref().map_or(false, |e| is_relevant_expr(e))
}

fn is_relevant_expr(expr: &Expr) -> bool {
    match expr.node {
        ExprBlock(ref block) => is_relevant_block(block),
        ExprRet(Some(ref e)) => is_relevant_expr(e),
        ExprRet(None) | ExprBreak(_) => false,
        ExprCall(ref path_expr, _) => {
            if let ExprPath(_, ref path) = path_expr.node {
                !match_path(path, &paths::BEGIN_PANIC)
            } else {
                true
            }
        }
        _ => true,
    }
}

fn check_attrs(cx: &LateContext, span: Span, name: &Name, attrs: &[Attribute]) {
    if in_macro(cx, span) {
        return;
    }

    for attr in attrs {
        if let MetaItemKind::List(ref inline, ref values) = attr.node.value.node {
            if values.len() != 1 || inline != &"inline" {
                continue;
            }
            if let MetaItemKind::Word(ref always) = values[0].node {
                if always != &"always" {
                    continue;
                }
                span_lint(cx,
                          INLINE_ALWAYS,
                          attr.span,
                          &format!("you have declared `#[inline(always)]` on `{}`. This is usually a bad idea",
                                   name));
            }
        }
    }
}

fn check_semver(cx: &LateContext, span: Span, lit: &Lit) {
    if let LitKind::Str(ref is, _) = lit.node {
        if Version::parse(&*is).is_ok() {
            return;
        }
    }
    span_lint(cx,
              DEPRECATED_SEMVER,
              span,
              "the since field must contain a semver-compliant version");
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use std::thread;
use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::io;
use std::io::Cursor;
use std::iter::repeat;

use solicit::http::server::ServerConnection;
use solicit::http::server::StreamFactory;
use solicit::http::server::ServerSession;
use solicit::http::transport::TransportStream;
use solicit::http::transport::TransportReceiveFrame;
use solicit::http::session::Server;
use solicit::http::session::SessionState;
use solicit::http::session::DefaultSessionState;
use solicit::http::session::DefaultStream;
use solicit::http::session::StreamState;
use solicit::http::session::StreamDataChunk;
use solicit::http::session::StreamDataError;
use solicit::http::session::Stream as solicit_Stream;
use solicit::http::connection::HttpConnection;
use solicit::http::connection::SendStatus;
use solicit::http::connection::HttpFrame;
use solicit::http::frame::RawFrame;
use solicit::http::HttpScheme;
use solicit::http::StreamId;
use solicit::http::Header;
use solicit::http::HttpResult;
use solicit::http::Response;
use solicit::http::HttpError;

use futures::done;
use futures::Future;
use futures::stream;
use futures::stream::Stream;
use futures_io::IoFuture;
use futures_io::TaskIo;
use futures_io::TaskIoRead;
use futures_io::TaskIoWrite;
use futures_mio::Loop;
use futures_mio::TcpStream;

use method::*;
use result::*;
use futures_grpc::*;
use futures_misc::*;
use solicit_async::*;
use solicit_misc::*;
use grpc::*;

pub trait MethodHandlerAsync<Req, Resp> {
    fn handle(&self, req: Req) -> GrpcFuture<Resp>;
}

pub struct MethodHandlerAsyncEcho;

impl<A : Send + 'static> MethodHandlerAsync<A, A> for MethodHandlerAsyncEcho {
    fn handle(&self, req: A) -> GrpcFuture<A> {
        println!("handle echo");
        done(Ok(req)).boxed()
    }
}

pub struct MethodHandlerAsyncFn<F> {
    f: F
}

impl<F> MethodHandlerAsyncFn<F> {
    pub fn new<Req, Resp>(f: F)
        -> Self
        where F : Fn(Req) -> GrpcFuture<Resp>
    {
        MethodHandlerAsyncFn {
            f: f,
        }
    }
}

impl<Req, Resp, F> MethodHandlerAsync<Req, Resp> for MethodHandlerAsyncFn<F>
    where
        Resp : Send + 'static,
        F : Fn(Req) -> GrpcFuture<Resp>,
{
    fn handle(&self, req: Req) -> GrpcFuture<Resp> {
        (self.f)(req)
    }
}

trait MethodHandlerDispatchAsync {
    fn on_message(&self, message: &[u8]) -> GrpcFuture<Vec<u8>>;
}

struct MethodHandlerDispatchAsyncImpl<Req, Resp> {
    desc: Arc<MethodDescriptor<Req, Resp>>,
    method_handler: Box<MethodHandlerAsync<Req, Resp> + Sync + Send>,
}

impl<Req, Resp> MethodHandlerDispatchAsync for MethodHandlerDispatchAsyncImpl<Req, Resp>
    where
        Req : Send + 'static,
        Resp : Send + 'static,
{
    fn on_message(&self, message: &[u8]) -> GrpcFuture<Vec<u8>> {
        let req = self.desc.req_marshaller.read(message);
        let resp = self.method_handler.handle(req);
        let desc_copy = self.desc.clone();
        resp
            .map(move |resp| desc_copy.resp_marshaller.write(&resp))
            .boxed()
    }
}

pub struct ServerMethodAsync {
    name: String,
    dispatch: Box<MethodHandlerDispatchAsync + Sync + Send>,
}

impl ServerMethodAsync {
    pub fn new<Req, Resp, H>(method: MethodDescriptor<Req, Resp>, handler: H) -> ServerMethodAsync
        where
            Req : Send + 'static,
            Resp : Send + 'static,
            H : MethodHandlerAsync<Req, Resp> + 'static + Sync + Send,
    {
        ServerMethodAsync {
            name: method.name.clone(),
            dispatch: Box::new(MethodHandlerDispatchAsyncImpl {
                desc: Arc::new(method),
                method_handler: Box::new(handler),
            }),
        }
    }
}

pub struct ServerServiceDefinitionAsync {
    methods: Vec<ServerMethodAsync>,
}

impl ServerServiceDefinitionAsync {
    pub fn new(mut methods: Vec<ServerMethodAsync>) -> ServerServiceDefinitionAsync {
        ServerServiceDefinitionAsync {
            methods: methods,
        }
    }

    pub fn find_method(&self, name: &str) -> &ServerMethodAsync {
        self.methods.iter()
            .filter(|m| m.name == name)
            .next()
            .expect(&format!("unknown method: {}", name))
    }

    pub fn handle_method(&self, name: &str, message: &[u8]) -> GrpcFuture<Vec<u8>> {
        self.find_method(name).dispatch.on_message(message)
    }
}

pub struct GrpcServerAsync {

}

impl GrpcServerAsync {
    pub fn new(port: u16, service_definition: ServerServiceDefinitionAsync) -> GrpcServerAsync {
        let listen_addr = ("::", port).to_socket_addrs().unwrap().next().unwrap();

        thread::spawn(move || {
            run_server_event_loop(listen_addr, service_definition);
        });

        GrpcServerAsync {
        }
    }
}

struct GrpcHttp2ServerStream {
    stream: DefaultStream,

    service_definition: Arc<ServerServiceDefinitionAsync>,
    path: String,
}

impl GrpcHttp2ServerStream {
    fn with_id(stream_id: StreamId, service_definition: Arc<ServerServiceDefinitionAsync>) -> Self {
        println!("new stream {}", stream_id);
        GrpcHttp2ServerStream {
            stream: DefaultStream::new(),
            service_definition: service_definition,
            path: String::new(),
        }
    }
}

impl solicit_Stream for GrpcHttp2ServerStream {
    fn set_headers(&mut self, headers: Vec<Header>) {
        self.stream.set_headers(headers)
    }

    fn new_data_chunk(&mut self, data: &[u8]) {
        self.stream.new_data_chunk(data)
    }

    fn set_state(&mut self, state: StreamState) {
        self.stream.set_state(state)
    }

    fn get_data_chunk(&mut self, buf: &mut [u8]) -> Result<StreamDataChunk, StreamDataError> {
        self.stream.get_data_chunk(buf)
    }

    fn state(&self) -> StreamState {
        self.stream.state()
    }
}

struct GrpcStreamFactory {
    service_definition: Arc<ServerServiceDefinitionAsync>,
}

impl StreamFactory for GrpcStreamFactory {
	type Stream = GrpcHttp2ServerStream;

	fn create(&mut self, id: StreamId) -> GrpcHttp2ServerStream {
		GrpcHttp2ServerStream::with_id(id, self.service_definition.clone())
	}
}

struct ServerSharedState {
    conn: ServerConnection<GrpcStreamFactory, DefaultSessionState<Server, GrpcHttp2ServerStream>>,
}

fn run_read(
    read: TaskIoRead<TcpStream>,
    shared: TaskDataMutex<ServerSharedState>)
        -> GrpcFuture<()>
{
    let stream = stream::iter(repeat(()).map(|x| Ok(x)));

    let future = stream.fold((read, shared), |(read, shared), _| {
        recv_raw_frame(read)
            .map(|(read, raw_frame)| {
                println!("received frame");

                shared.with(|shared| {

                    let mut send = VecSendFrame(Vec::new());

                    shared.conn.handle_next_frame(
                        &mut OnceReceiveFrame::new(raw_frame),
                        &mut send);

                    // TODO: process send
                });
                (read, shared)
            })
    });

    future
        .map(|_| ())
        .boxed()
}

fn run_write(
    write: TaskIoWrite<TcpStream>,
    shared: TaskDataMutex<ServerSharedState>)
        -> GrpcFuture<()>
{
    done(Ok(())).boxed()
}

fn run_connection(
    socket: TcpStream,
    peer_addr: SocketAddr,
    service_defintion: Arc<ServerServiceDefinitionAsync>)
        -> IoFuture<()>
{
    println!("accepted connection from {}", peer_addr);
    let handshake = server_handshake(socket)
        .map_err(|e| e.into());

    let run = handshake.and_then(|conn| {
        let (read, write) = TaskIo::new(conn).split();

        let shared_for_read = TaskDataMutex::new(ServerSharedState {
            conn: ServerConnection::with_connection(
                HttpConnection::new(HttpScheme::Http),
                DefaultSessionState::<Server, _>::new(),
                GrpcStreamFactory {
                    // TODO
                    service_definition: service_defintion,
                }
            ),
        });
        let shared_for_write = shared_for_read.clone();

        run_read(read, shared_for_read)
            .join(run_write(write, shared_for_write))
    });

    let bye = run
        .then(move |x| {
            println!("closing connection from {}: {:?}", peer_addr, x);
            x
        });

    bye
        .map(|_| ())
        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
        .boxed()
}

fn run_server_event_loop(listen_addr: SocketAddr, service_definition: ServerServiceDefinitionAsync) {
    let mut lp = Loop::new().unwrap();

    let listen = lp.handle().tcp_listen(&listen_addr);

    let service_definition_stream = stream_repeat(Arc::new(service_definition));

    let done = listen.and_then(|socket| {
        socket.incoming().zip(service_definition_stream).for_each(|((socket, peer_addr), service_definition)| {
            run_connection(socket, peer_addr, service_definition).forget();
            Ok(())
        })
    });

    lp.run(done).unwrap();
}
<commit_msg>backup<commit_after>use std::sync::Arc;
use std::thread;
use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::io;
use std::io::Cursor;
use std::iter::repeat;

use solicit::http::server::ServerConnection;
use solicit::http::server::StreamFactory;
use solicit::http::server::ServerSession;
use solicit::http::transport::TransportStream;
use solicit::http::transport::TransportReceiveFrame;
use solicit::http::session::Server;
use solicit::http::session::SessionState;
use solicit::http::session::DefaultSessionState;
use solicit::http::session::DefaultStream;
use solicit::http::session::StreamState;
use solicit::http::session::StreamDataChunk;
use solicit::http::session::StreamDataError;
use solicit::http::session::Stream as solicit_Stream;
use solicit::http::connection::HttpConnection;
use solicit::http::connection::SendStatus;
use solicit::http::connection::HttpFrame;
use solicit::http::frame::RawFrame;
use solicit::http::HttpScheme;
use solicit::http::StreamId;
use solicit::http::Header;
use solicit::http::HttpResult;
use solicit::http::Response;
use solicit::http::HttpError;

use futures::done;
use futures::Future;
use futures::stream;
use futures::stream::Stream;
use futures_io::IoFuture;
use futures_io::TaskIo;
use futures_io::TaskIoRead;
use futures_io::TaskIoWrite;
use futures_mio::Loop;
use futures_mio::TcpStream;

use method::*;
use result::*;
use futures_grpc::*;
use futures_misc::*;
use solicit_async::*;
use solicit_misc::*;
use grpc::*;

pub trait MethodHandlerAsync<Req, Resp> {
    fn handle(&self, req: Req) -> GrpcFuture<Resp>;
}

pub struct MethodHandlerAsyncEcho;

impl<A : Send + 'static> MethodHandlerAsync<A, A> for MethodHandlerAsyncEcho {
    fn handle(&self, req: A) -> GrpcFuture<A> {
        println!("handle echo");
        done(Ok(req)).boxed()
    }
}

pub struct MethodHandlerAsyncFn<F> {
    f: F
}

impl<F> MethodHandlerAsyncFn<F> {
    pub fn new<Req, Resp>(f: F)
        -> Self
        where F : Fn(Req) -> GrpcFuture<Resp>
    {
        MethodHandlerAsyncFn {
            f: f,
        }
    }
}

impl<Req, Resp, F> MethodHandlerAsync<Req, Resp> for MethodHandlerAsyncFn<F>
    where
        Resp : Send + 'static,
        F : Fn(Req) -> GrpcFuture<Resp>,
{
    fn handle(&self, req: Req) -> GrpcFuture<Resp> {
        (self.f)(req)
    }
}

trait MethodHandlerDispatchAsync {
    fn on_message(&self, message: &[u8]) -> GrpcFuture<Vec<u8>>;
}

struct MethodHandlerDispatchAsyncImpl<Req, Resp> {
    desc: Arc<MethodDescriptor<Req, Resp>>,
    method_handler: Box<MethodHandlerAsync<Req, Resp> + Sync + Send>,
}

impl<Req, Resp> MethodHandlerDispatchAsync for MethodHandlerDispatchAsyncImpl<Req, Resp>
    where
        Req : Send + 'static,
        Resp : Send + 'static,
{
    fn on_message(&self, message: &[u8]) -> GrpcFuture<Vec<u8>> {
        let req = self.desc.req_marshaller.read(message);
        let resp = self.method_handler.handle(req);
        let desc_copy = self.desc.clone();
        resp
            .map(move |resp| desc_copy.resp_marshaller.write(&resp))
            .boxed()
    }
}

pub struct ServerMethodAsync {
    name: String,
    dispatch: Box<MethodHandlerDispatchAsync + Sync + Send>,
}

impl ServerMethodAsync {
    pub fn new<Req, Resp, H>(method: MethodDescriptor<Req, Resp>, handler: H) -> ServerMethodAsync
        where
            Req : Send + 'static,
            Resp : Send + 'static,
            H : MethodHandlerAsync<Req, Resp> + 'static + Sync + Send,
    {
        ServerMethodAsync {
            name: method.name.clone(),
            dispatch: Box::new(MethodHandlerDispatchAsyncImpl {
                desc: Arc::new(method),
                method_handler: Box::new(handler),
            }),
        }
    }
}

pub struct ServerServiceDefinitionAsync {
    methods: Vec<ServerMethodAsync>,
}

impl ServerServiceDefinitionAsync {
    pub fn new(mut methods: Vec<ServerMethodAsync>) -> ServerServiceDefinitionAsync {
        ServerServiceDefinitionAsync {
            methods: methods,
        }
    }

    pub fn find_method(&self, name: &str) -> &ServerMethodAsync {
        self.methods.iter()
            .filter(|m| m.name == name)
            .next()
            .expect(&format!("unknown method: {}", name))
    }

    pub fn handle_method(&self, name: &str, message: &[u8]) -> GrpcFuture<Vec<u8>> {
        self.find_method(name).dispatch.on_message(message)
    }
}

pub struct GrpcServerAsync {

}

impl GrpcServerAsync {
    pub fn new(port: u16, service_definition: ServerServiceDefinitionAsync) -> GrpcServerAsync {
        let listen_addr = ("::", port).to_socket_addrs().unwrap().next().unwrap();

        thread::spawn(move || {
            run_server_event_loop(listen_addr, service_definition);
        });

        GrpcServerAsync {
        }
    }
}

struct GrpcHttp2ServerStream {
    stream: DefaultStream,

    service_definition: Arc<ServerServiceDefinitionAsync>,
    path: String,
}

impl GrpcHttp2ServerStream {
    fn with_id(stream_id: StreamId, service_definition: Arc<ServerServiceDefinitionAsync>) -> Self {
        println!("new stream {}", stream_id);
        GrpcHttp2ServerStream {
            stream: DefaultStream::new(),
            service_definition: service_definition,
            path: String::new(),
        }
    }
}

impl solicit_Stream for GrpcHttp2ServerStream {
    fn set_headers(&mut self, headers: Vec<Header>) {
        for h in &headers {
            if h.name() == b":path" {
                self.path = String::from_utf8(h.value().to_owned()).unwrap();
            }
        }

        self.stream.set_headers(headers)
    }

    fn new_data_chunk(&mut self, data: &[u8]) {
        self.stream.new_data_chunk(data)
    }

    fn set_state(&mut self, state: StreamState) {
        println!("set state {:?}", state);
        if state == StreamState::HalfClosedRemote {
            let message = parse_grpc_frame_completely(&self.stream.body).unwrap();
            self.service_definition.handle_method(&self.path, message)
                .forget();
        }
        self.stream.set_state(state)
    }

    fn get_data_chunk(&mut self, buf: &mut [u8]) -> Result<StreamDataChunk, StreamDataError> {
        self.stream.get_data_chunk(buf)
    }

    fn state(&self) -> StreamState {
        self.stream.state()
    }
}

struct GrpcStreamFactory {
    service_definition: Arc<ServerServiceDefinitionAsync>,
}

impl StreamFactory for GrpcStreamFactory {
	type Stream = GrpcHttp2ServerStream;

	fn create(&mut self, id: StreamId) -> GrpcHttp2ServerStream {
		GrpcHttp2ServerStream::with_id(id, self.service_definition.clone())
	}
}

struct ServerSharedState {
    conn: ServerConnection<GrpcStreamFactory, DefaultSessionState<Server, GrpcHttp2ServerStream>>,
}

fn run_read(
    read: TaskIoRead<TcpStream>,
    shared: TaskDataMutex<ServerSharedState>)
        -> GrpcFuture<()>
{
    let stream = stream::iter(repeat(()).map(|x| Ok(x)));

    let future = stream.fold((read, shared), |(read, shared), _| {
        recv_raw_frame(read)
            .map(|(read, raw_frame)| {
                println!("received frame");

                shared.with(|shared| {

                    let mut send = VecSendFrame(Vec::new());

                    shared.conn.handle_next_frame(
                        &mut OnceReceiveFrame::new(raw_frame),
                        &mut send);

                    // TODO: process send
                });
                (read, shared)
            })
    });

    future
        .map(|_| ())
        .boxed()
}

fn run_write(
    write: TaskIoWrite<TcpStream>,
    shared: TaskDataMutex<ServerSharedState>)
        -> GrpcFuture<()>
{
    done(Ok(())).boxed()
}

fn run_connection(
    socket: TcpStream,
    peer_addr: SocketAddr,
    service_defintion: Arc<ServerServiceDefinitionAsync>)
        -> IoFuture<()>
{
    println!("accepted connection from {}", peer_addr);
    let handshake = server_handshake(socket)
        .map_err(|e| e.into());

    let run = handshake.and_then(|conn| {
        let (read, write) = TaskIo::new(conn).split();

        let shared_for_read = TaskDataMutex::new(ServerSharedState {
            conn: ServerConnection::with_connection(
                HttpConnection::new(HttpScheme::Http),
                DefaultSessionState::<Server, _>::new(),
                GrpcStreamFactory {
                    // TODO
                    service_definition: service_defintion,
                }
            ),
        });
        let shared_for_write = shared_for_read.clone();

        run_read(read, shared_for_read)
            .join(run_write(write, shared_for_write))
    });

    let bye = run
        .then(move |x| {
            println!("closing connection from {}: {:?}", peer_addr, x);
            x
        });

    bye
        .map(|_| ())
        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
        .boxed()
}

fn run_server_event_loop(listen_addr: SocketAddr, service_definition: ServerServiceDefinitionAsync) {
    let mut lp = Loop::new().unwrap();

    let listen = lp.handle().tcp_listen(&listen_addr);

    let service_definition_stream = stream_repeat(Arc::new(service_definition));

    let done = listen.and_then(|socket| {
        socket.incoming().zip(service_definition_stream).for_each(|((socket, peer_addr), service_definition)| {
            run_connection(socket, peer_addr, service_definition).forget();
            Ok(())
        })
    });

    lp.run(done).unwrap();
}
<|endoftext|>