<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    network_interface::OnchainDiscoveryNetworkSender,
    storage_query_discovery_set_async,
    types::{
        DiscoveryInfoInternal, DiscoverySetInternal, QueryDiscoverySetRequest,
        QueryDiscoverySetResponse,
    },
};
use anyhow::{Context as _, Result};
use futures::{
    channel::mpsc,
    future::{Future, FutureExt},
    sink::SinkExt,
    stream::{FusedStream, Stream, StreamExt},
};
use libra_config::config::RoleType;
use libra_logger::prelude::*;
use libra_types::{
    on_chain_config::ValidatorSet,
    trusted_state::{TrustedState, TrustedStateChange},
    waypoint::Waypoint,
    PeerId,
};
use network::{
    connectivity_manager::{ConnectivityRequest, DiscoverySource},
    peer_manager::{conn_notifs_channel, ConnectionNotification},
};
use option_future::OptionFuture;
use rand::{rngs::SmallRng, Rng, SeedableRng};
use std::{collections::HashSet, mem, sync::Arc, time::Duration};
use storage_interface::DbReader;

/// Actor for querying various sources (remote peers, local storage) for the
/// latest discovery set and notifying the `ConnectivityManager` of updates.
pub struct OnchainDiscovery<TTicker> {
    /// Our node's PeerId.
    peer_id: PeerId,
    /// Our node's role (validator || fullnode).
    role: RoleType,
    /// The current trusted state, which keeps track of the latest verified
    /// ledger version and validator set. This version can be ahead of other
    /// components (e.g., state sync) since we only rely on syncing epoch change
    /// LedgerInfo's and don't need to do tx replay.
    trusted_state: TrustedState,
    /// An internal representation of the most recent discovery set.
    latest_discovery_set: DiscoverySetInternal,
    /// The set of peers we're connected to.
    connected_peers: HashSet<PeerId>,
    /// A channel to send requests to the network instance.
    network_tx: OnchainDiscoveryNetworkSender,
    /// A channel to send requests to the connectivity manager
    conn_mgr_reqs_tx: channel::Sender<ConnectivityRequest>,
    /// A channel to receive connection updates from the network.
    conn_notifs_rx: conn_notifs_channel::Receiver,
    /// internal gRPC client to send read requests to Libra Storage.
    libra_db: Arc<dyn DbReader>,
    /// Ticker to periodically query our peers for the latest discovery set.
    // TODO(philiphayes): once we do a bunch of initial queries on startup, we
    // can probably reduce the frequency.
    peer_query_ticker: TTicker,
    /// Ticker to query our own storage for the latest discovery set.
    // TODO(philiphayes): we should only need to query our own storage on startup
    // and after that rely on state-sync's reconfig notifications.
    storage_query_ticker: TTicker,
    /// The timeout for the whole rpc request when querying one of our peers for
    /// their latest discovery set.
    outbound_rpc_timeout: Duration,
    /// Random number generator. Used to sample the next peer to query.
    rng: SmallRng,
    /// A local counter incremented on receiving an incoming message. Printing this in debugging
    /// allows for easy debugging.
    event_id: u32,
}

impl<TTicker> OnchainDiscovery<TTicker>
where
    TTicker: Stream + FusedStream + Unpin,
{
    pub fn new(
        self_peer_id: PeerId,
        role: RoleType,
        waypoint: Waypoint,
        network_tx: OnchainDiscoveryNetworkSender,
        conn_mgr_reqs_tx: channel::Sender<ConnectivityRequest>,
        conn_notifs_rx: conn_notifs_channel::Receiver,
        libra_db: Arc<dyn DbReader>,
        peer_query_ticker: TTicker,
        storage_query_ticker: TTicker,
        outbound_rpc_timeout: Duration,
    ) -> Self {
        let trusted_state = waypoint.into();

        Self {
            peer_id: self_peer_id,
            role,
            trusted_state,
            latest_discovery_set: DiscoverySetInternal::empty(),
            connected_peers: HashSet::new(),
            network_tx,
            conn_mgr_reqs_tx,
            conn_notifs_rx,
            libra_db,
            peer_query_ticker,
            storage_query_ticker,
            outbound_rpc_timeout,
            rng: SmallRng::from_entropy(),
            event_id: 0,
        }
    }

    pub async fn start(mut self) {
        // A slot that can hold at-most-one pending outbound discovery set query
        // to another peer.
        let mut pending_outbound_rpc = OptionFuture::new(None);

        // On startup, we want to first query our local storage state to see if
        // we have an up-to-date discovery set to connect with.
        let f_query_storage = self.handle_storage_query_tick().boxed();
        let mut pending_storage_query = OptionFuture::new(Some(f_query_storage));

        debug!("starting onchain discovery");
        loop {
            self.event_id = self.event_id.wrapping_add(1);
            futures::select! {
                notif = self.conn_notifs_rx.select_next_some() => {
                    trace!("event id: {}, type: ConnectionNotification", self.event_id);
                    self.handle_connection_notif(notif);
                },
                _ = self.peer_query_ticker.select_next_some() => {
                    trace!("event id: {}, type: PeerQueryTick", self.event_id);
                    pending_outbound_rpc
                        .or_insert_with(|| {
                            self.handle_peer_query_tick()
                                .map(|f_query| f_query.boxed())
                        });
                },
                query_res = pending_outbound_rpc => {
                    trace!("event id: {}, type: QueryResponse", self.event_id);
                    self.handle_outbound_query_res(query_res).await;
                },
                _ = self.storage_query_ticker.select_next_some() => {
                    trace!("event id: {}, type: StorageQueryTick", self.event_id);
                    pending_storage_query
                        .or_insert_with(|| Some(self.handle_storage_query_tick().boxed()));
                },
                query_res = pending_storage_query => {
                    trace!("event id: {}, type: QueryResponse", self.event_id);
                    self.handle_outbound_query_res(query_res).await;
                }
                complete => {
                    crit!("onchain discovery terminated");
                    break;
                }
            }
        }
    }

    fn handle_connection_notif(&mut self, notif: ConnectionNotification) {
        match notif {
            ConnectionNotification::NewPeer(peer_id, _addr) => {
                trace!("connected to new peer: {}", peer_id.short_str());
                // Add peer to connected peer list.
                self.connected_peers.insert(peer_id);
            }
            ConnectionNotification::LostPeer(peer_id, _addr, _reason) => {
                trace!("disconnected from peer: {}", peer_id.short_str());
                // Remove peer from connected peer list.
                self.connected_peers.remove(&peer_id);
            }
        }
    }

    fn handle_peer_query_tick(
        &mut self,
    ) -> Option<
        impl Future<
            Output = Result<(
                PeerId,
                QueryDiscoverySetRequest,
                Box<QueryDiscoverySetResponse>,
            )>,
        >,
    > {
        self.sample_peer().map(|peer_id| {
            let req_msg = QueryDiscoverySetRequest {
                known_version: self.trusted_state.latest_version(),
            };

            let peer_id_short = peer_id.short_str();

            trace!(
                "handle_peer_query_tick: querying peer: {}, trusted version: {}",
                peer_id_short,
                req_msg.known_version,
            );

            peer_query_discovery_set(
                peer_id,
                req_msg,
                self.outbound_rpc_timeout,
                self.network_tx.clone(),
            )
        })
    }

    async fn handle_outbound_query_res(
        &mut self,
        query_res: Result<(
            PeerId,
            QueryDiscoverySetRequest,
            Box<QueryDiscoverySetResponse>,
        )>,
    ) {
        match query_res {
            Ok((peer_id, req_msg, res_msg)) => {
                debug!(
                    "received query response: peer: {}, their version: {}",
                    peer_id.short_str(),
                    res_msg.latest_li.ledger_info().version(),
                );
                self.handle_query_response(peer_id, req_msg, res_msg).await;
            }
            Err(err) => warn!("query to remote peer failed: {:?}", err),
        }
    }

    fn handle_storage_query_tick(
        &self,
    ) -> impl Future<
        Output = Result<(
            PeerId,
            QueryDiscoverySetRequest,
            Box<QueryDiscoverySetResponse>,
        )>,
    > {
        let trusted_version = self.trusted_state.latest_version();
        trace!(
            "handle_storage_query_tick: querying self storage, trusted version: {}",
            trusted_version,
        );
        let req_msg = QueryDiscoverySetRequest {
            known_version: trusted_version,
        };
        let self_peer_id = self.peer_id;
        let libra_db = Arc::clone(&self.libra_db);
        storage_query_discovery_set_async(libra_db, req_msg).map(move |res| {
            // add peer_id context
            res.map(|(req_msg, res_msg)| (self_peer_id, req_msg, res_msg))
        })
    }

    async fn handle_query_response(
        &mut self,
        peer_id: PeerId,
        req_msg: QueryDiscoverySetRequest,
        res_msg: Box<QueryDiscoverySetResponse>,
    ) {
        let (trusted_state_change, opt_discovery_set) =
            match res_msg.verify_and_ratchet(&req_msg, &self.trusted_state) {
                Ok(res) => res,
                Err(err) => {
                    warn!(
                        "invalid query response: peer: {}, request_version: {}, err: {:?}",
                        peer_id.short_str(),
                        req_msg.known_version,
                        err
                    );
                    return;
                }
            };

        match trusted_state_change {
            TrustedStateChange::Epoch {
                new_state,
                latest_epoch_change_li,
            } => {
                info!(
                    "successfully ratcheted to new epoch: \
                     peer: {}, epoch: {}, version: {}",
                    peer_id.short_str(),
                    latest_epoch_change_li.ledger_info().epoch(),
                    new_state.latest_version(),
                );
                self.trusted_state = new_state;

                let discovery_set = opt_discovery_set.expect(
                    "We're guaranteed Some(DiscoverySet) from QueryDiscoverySetResponse::verify_and_ratchet"
                );
                self.handle_new_discovery_set_event(discovery_set).await;
            }
            TrustedStateChange::Version { new_state } => {
                debug!(
                    "successfully ratcheted to new version: \
                     peer: {}, version: {}",
                    peer_id.short_str(),
                    new_state.latest_version(),
                );
                self.trusted_state = new_state;
            }
            TrustedStateChange::NoChange => (),
        };
    }

    async fn send_connectivity_request(
        &mut self,
        req: ConnectivityRequest,
    ) -> Result<(), mpsc::SendError> {
        self.conn_mgr_reqs_tx.send(req).await
    }

    async fn handle_new_discovery_set_event(&mut self, validator_set: ValidatorSet) {
        let latest_discovery_set =
            DiscoverySetInternal::from_validator_set(self.role, validator_set);

        let prev_discovery_set =
            mem::replace(&mut self.latest_discovery_set, latest_discovery_set.clone());

        // TODO(philiphayes): ConnectivityManager supports multiple identity pubkeys per
        // peer (will accept connections from any in set, and do { pubkeys } x { addrs }
        // when attempting to dial).

        // Compare the latest and previous discovery sets to determine if
        // we need to update the connectivity manager that a peer is
        // advertising new addresses.
        if self.latest_discovery_set != prev_discovery_set {
            let update_addr_reqs = self
                .latest_discovery_set
                .0
                .iter()
                .map(|(peer_id, DiscoveryInfoInternal(_id_pubkey, addrs))| {
                    (*peer_id, addrs.clone())
                })
                .collect();
            self.send_connectivity_request(ConnectivityRequest::UpdateAddresses(
                DiscoverySource::OnChain,
                update_addr_reqs,
            ))
            .await
            .unwrap();
        }
    }

    /// Sample iid a PeerId of one of the connected peers.
    fn sample_peer(&mut self) -> Option<PeerId> {
        let num_peers = self.connected_peers.len();
        if num_peers > 0 {
            // The particular index is not important, just the probability of
            // sampling any of the current peers is equal.
            let idx = self.rng.gen_range(0, num_peers);
            self.connected_peers.iter().nth(idx).copied()
        } else {
            None
        }
    }
}

/// Query a remote peer for their latest discovery set and a epoch change
/// proof.
async fn peer_query_discovery_set(
    peer_id: PeerId,
    req_msg: QueryDiscoverySetRequest,
    outbound_rpc_timeout: Duration,
    mut network_tx: OnchainDiscoveryNetworkSender,
) -> Result<(
    PeerId,
    QueryDiscoverySetRequest,
    Box<QueryDiscoverySetResponse>,
)> {
    let our_latest_version = req_msg.known_version;
    let res_msg = network_tx
        .query_discovery_set(peer_id, req_msg.clone(), outbound_rpc_timeout)
        .await
        .with_context(|| {
            format!(
                "failed to query peer discovery set: peer: {}, latest_version: {}",
                peer_id.short_str(),
                our_latest_version,
            )
        })?;

    Ok((peer_id, req_msg, res_msg))
}
<commit_msg>[onchaindiscovery] inline a call to send<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    network_interface::OnchainDiscoveryNetworkSender,
    storage_query_discovery_set_async,
    types::{
        DiscoveryInfoInternal, DiscoverySetInternal, QueryDiscoverySetRequest,
        QueryDiscoverySetResponse,
    },
};
use anyhow::{Context as _, Result};
use futures::{
    future::{Future, FutureExt},
    sink::SinkExt,
    stream::{FusedStream, Stream, StreamExt},
};
use libra_config::config::RoleType;
use libra_logger::prelude::*;
use libra_types::{
    on_chain_config::ValidatorSet,
    trusted_state::{TrustedState, TrustedStateChange},
    waypoint::Waypoint,
    PeerId,
};
use network::{
    connectivity_manager::{ConnectivityRequest, DiscoverySource},
    peer_manager::{conn_notifs_channel, ConnectionNotification},
};
use option_future::OptionFuture;
use rand::{rngs::SmallRng, Rng, SeedableRng};
use std::{collections::HashSet, mem, sync::Arc, time::Duration};
use storage_interface::DbReader;

/// Actor for querying various sources (remote peers, local storage) for the
/// latest discovery set and notifying the `ConnectivityManager` of updates.
pub struct OnchainDiscovery<TTicker> {
    /// Our node's PeerId.
    peer_id: PeerId,
    /// Our node's role (validator || fullnode).
    role: RoleType,
    /// The current trusted state, which keeps track of the latest verified
    /// ledger version and validator set. This version can be ahead of other
    /// components (e.g., state sync) since we only rely on syncing epoch change
    /// LedgerInfo's and don't need to do tx replay.
    trusted_state: TrustedState,
    /// An internal representation of the most recent discovery set.
    latest_discovery_set: DiscoverySetInternal,
    /// The set of peers we're connected to.
    connected_peers: HashSet<PeerId>,
    /// A channel to send requests to the network instance.
    network_tx: OnchainDiscoveryNetworkSender,
    /// A channel to send requests to the connectivity manager
    conn_mgr_reqs_tx: channel::Sender<ConnectivityRequest>,
    /// A channel to receive connection updates from the network.
    conn_notifs_rx: conn_notifs_channel::Receiver,
    /// internal gRPC client to send read requests to Libra Storage.
    libra_db: Arc<dyn DbReader>,
    /// Ticker to periodically query our peers for the latest discovery set.
    // TODO(philiphayes): once we do a bunch of initial queries on startup, we
    // can probably reduce the frequency.
    peer_query_ticker: TTicker,
    /// Ticker to query our own storage for the latest discovery set.
    // TODO(philiphayes): we should only need to query our own storage on startup
    // and after that rely on state-sync's reconfig notifications.
    storage_query_ticker: TTicker,
    /// The timeout for the whole rpc request when querying one of our peers for
    /// their latest discovery set.
    outbound_rpc_timeout: Duration,
    /// Random number generator. Used to sample the next peer to query.
    rng: SmallRng,
    /// A local counter incremented on receiving an incoming message. Printing this in debugging
    /// allows for easy debugging.
    event_id: u32,
}

impl<TTicker> OnchainDiscovery<TTicker>
where
    TTicker: Stream + FusedStream + Unpin,
{
    pub fn new(
        self_peer_id: PeerId,
        role: RoleType,
        waypoint: Waypoint,
        network_tx: OnchainDiscoveryNetworkSender,
        conn_mgr_reqs_tx: channel::Sender<ConnectivityRequest>,
        conn_notifs_rx: conn_notifs_channel::Receiver,
        libra_db: Arc<dyn DbReader>,
        peer_query_ticker: TTicker,
        storage_query_ticker: TTicker,
        outbound_rpc_timeout: Duration,
    ) -> Self {
        let trusted_state = waypoint.into();

        Self {
            peer_id: self_peer_id,
            role,
            trusted_state,
            latest_discovery_set: DiscoverySetInternal::empty(),
            connected_peers: HashSet::new(),
            network_tx,
            conn_mgr_reqs_tx,
            conn_notifs_rx,
            libra_db,
            peer_query_ticker,
            storage_query_ticker,
            outbound_rpc_timeout,
            rng: SmallRng::from_entropy(),
            event_id: 0,
        }
    }

    pub async fn start(mut self) {
        // A slot that can hold at-most-one pending outbound discovery set query
        // to another peer.
        let mut pending_outbound_rpc = OptionFuture::new(None);

        // On startup, we want to first query our local storage state to see if
        // we have an up-to-date discovery set to connect with.
        let f_query_storage = self.handle_storage_query_tick().boxed();
        let mut pending_storage_query = OptionFuture::new(Some(f_query_storage));

        debug!("starting onchain discovery");
        loop {
            self.event_id = self.event_id.wrapping_add(1);
            futures::select! {
                notif = self.conn_notifs_rx.select_next_some() => {
                    trace!("event id: {}, type: ConnectionNotification", self.event_id);
                    self.handle_connection_notif(notif);
                },
                _ = self.peer_query_ticker.select_next_some() => {
                    trace!("event id: {}, type: PeerQueryTick", self.event_id);
                    pending_outbound_rpc
                        .or_insert_with(|| {
                            self.handle_peer_query_tick()
                                .map(|f_query| f_query.boxed())
                        });
                },
                query_res = pending_outbound_rpc => {
                    trace!("event id: {}, type: QueryResponse", self.event_id);
                    self.handle_outbound_query_res(query_res).await;
                },
                _ = self.storage_query_ticker.select_next_some() => {
                    trace!("event id: {}, type: StorageQueryTick", self.event_id);
                    pending_storage_query
                        .or_insert_with(|| Some(self.handle_storage_query_tick().boxed()));
                },
                query_res = pending_storage_query => {
                    trace!("event id: {}, type: QueryResponse", self.event_id);
                    self.handle_outbound_query_res(query_res).await;
                }
                complete => {
                    crit!("onchain discovery terminated");
                    break;
                }
            }
        }
    }

    fn handle_connection_notif(&mut self, notif: ConnectionNotification) {
        match notif {
            ConnectionNotification::NewPeer(peer_id, _addr) => {
                trace!("connected to new peer: {}", peer_id.short_str());
                // Add peer to connected peer list.
                self.connected_peers.insert(peer_id);
            }
            ConnectionNotification::LostPeer(peer_id, _addr, _reason) => {
                trace!("disconnected from peer: {}", peer_id.short_str());
                // Remove peer from connected peer list.
                self.connected_peers.remove(&peer_id);
            }
        }
    }

    fn handle_peer_query_tick(
        &mut self,
    ) -> Option<
        impl Future<
            Output = Result<(
                PeerId,
                QueryDiscoverySetRequest,
                Box<QueryDiscoverySetResponse>,
            )>,
        >,
    > {
        self.sample_peer().map(|peer_id| {
            let req_msg = QueryDiscoverySetRequest {
                known_version: self.trusted_state.latest_version(),
            };

            let peer_id_short = peer_id.short_str();

            trace!(
                "handle_peer_query_tick: querying peer: {}, trusted version: {}",
                peer_id_short,
                req_msg.known_version,
            );

            peer_query_discovery_set(
                peer_id,
                req_msg,
                self.outbound_rpc_timeout,
                self.network_tx.clone(),
            )
        })
    }

    async fn handle_outbound_query_res(
        &mut self,
        query_res: Result<(
            PeerId,
            QueryDiscoverySetRequest,
            Box<QueryDiscoverySetResponse>,
        )>,
    ) {
        match query_res {
            Ok((peer_id, req_msg, res_msg)) => {
                debug!(
                    "received query response: peer: {}, their version: {}",
                    peer_id.short_str(),
                    res_msg.latest_li.ledger_info().version(),
                );
                self.handle_query_response(peer_id, req_msg, res_msg).await;
            }
            Err(err) => warn!("query to remote peer failed: {:?}", err),
        }
    }

    fn handle_storage_query_tick(
        &self,
    ) -> impl Future<
        Output = Result<(
            PeerId,
            QueryDiscoverySetRequest,
            Box<QueryDiscoverySetResponse>,
        )>,
    > {
        let trusted_version = self.trusted_state.latest_version();
        trace!(
            "handle_storage_query_tick: querying self storage, trusted version: {}",
            trusted_version,
        );
        let req_msg = QueryDiscoverySetRequest {
            known_version: trusted_version,
        };
        let self_peer_id = self.peer_id;
        let libra_db = Arc::clone(&self.libra_db);
        storage_query_discovery_set_async(libra_db, req_msg).map(move |res| {
            // add peer_id context
            res.map(|(req_msg, res_msg)| (self_peer_id, req_msg, res_msg))
        })
    }

    async fn handle_query_response(
        &mut self,
        peer_id: PeerId,
        req_msg: QueryDiscoverySetRequest,
        res_msg: Box<QueryDiscoverySetResponse>,
    ) {
        let (trusted_state_change, opt_discovery_set) =
            match res_msg.verify_and_ratchet(&req_msg, &self.trusted_state) {
                Ok(res) => res,
                Err(err) => {
                    warn!(
                        "invalid query response: peer: {}, request_version: {}, err: {:?}",
                        peer_id.short_str(),
                        req_msg.known_version,
                        err
                    );
                    return;
                }
            };

        match trusted_state_change {
            TrustedStateChange::Epoch {
                new_state,
                latest_epoch_change_li,
            } => {
                info!(
                    "successfully ratcheted to new epoch: \
                     peer: {}, epoch: {}, version: {}",
                    peer_id.short_str(),
                    latest_epoch_change_li.ledger_info().epoch(),
                    new_state.latest_version(),
                );
                self.trusted_state = new_state;

                let discovery_set = opt_discovery_set.expect(
                    "We're guaranteed Some(DiscoverySet) from QueryDiscoverySetResponse::verify_and_ratchet"
                );
                self.handle_new_discovery_set_event(discovery_set).await;
            }
            TrustedStateChange::Version { new_state } => {
                debug!(
                    "successfully ratcheted to new version: \
                     peer: {}, version: {}",
                    peer_id.short_str(),
                    new_state.latest_version(),
                );
                self.trusted_state = new_state;
            }
            TrustedStateChange::NoChange => (),
        };
    }

    async fn handle_new_discovery_set_event(&mut self, validator_set: ValidatorSet) {
        let latest_discovery_set =
            DiscoverySetInternal::from_validator_set(self.role, validator_set);

        let prev_discovery_set =
            mem::replace(&mut self.latest_discovery_set, latest_discovery_set.clone());

        // TODO(philiphayes): ConnectivityManager supports multiple identity pubkeys per
        // peer (will accept connections from any in set, and do { pubkeys } x { addrs }
        // when attempting to dial).

        // Compare the latest and previous discovery sets to determine if
        // we need to update the connectivity manager that a peer is
        // advertising new addresses.
        if self.latest_discovery_set != prev_discovery_set {
            let update_addr_reqs = self
                .latest_discovery_set
                .0
                .iter()
                .map(|(peer_id, DiscoveryInfoInternal(_id_pubkey, addrs))| {
                    (*peer_id, addrs.clone())
                })
                .collect();
            self.conn_mgr_reqs_tx
                .send(ConnectivityRequest::UpdateAddresses(
                    DiscoverySource::OnChain,
                    update_addr_reqs,
                ))
                .await
                .unwrap();
        }
    }

    /// Sample iid a PeerId of one of the connected peers.
    fn sample_peer(&mut self) -> Option<PeerId> {
        let num_peers = self.connected_peers.len();
        if num_peers > 0 {
            // The particular index is not important, just the probability of
            // sampling any of the current peers is equal.
            let idx = self.rng.gen_range(0, num_peers);
            self.connected_peers.iter().nth(idx).copied()
        } else {
            None
        }
    }
}

/// Query a remote peer for their latest discovery set and a epoch change
/// proof.
async fn peer_query_discovery_set(
    peer_id: PeerId,
    req_msg: QueryDiscoverySetRequest,
    outbound_rpc_timeout: Duration,
    mut network_tx: OnchainDiscoveryNetworkSender,
) -> Result<(
    PeerId,
    QueryDiscoverySetRequest,
    Box<QueryDiscoverySetResponse>,
)> {
    let our_latest_version = req_msg.known_version;
    let res_msg = network_tx
        .query_discovery_set(peer_id, req_msg.clone(), outbound_rpc_timeout)
        .await
        .with_context(|| {
            format!(
                "failed to query peer discovery set: peer: {}, latest_version: {}",
                peer_id.short_str(),
                our_latest_version,
            )
        })?;

    Ok((peer_id, req_msg, res_msg))
}
<|endoftext|>


<commit_before>extern crate time;
extern crate rand;
extern crate timely;
extern crate timely_sort;

use std::collections::HashMap;

use rand::{Rng, SeedableRng, StdRng};
use timely_sort::LSBRadixSorter as RadixSorter;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::dataflow::channels::pact::Exchange;
fn main() {

    // command-line args: numbers of nodes and edges in the random graph.
    let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();

    timely::execute_from_args(std::env::args().skip(3), move |root| {

        let index = root.index();
        let peers = root.peers();

        let seed: &[_] = &[1, 2, 3, root.index()];
        let mut rng: StdRng = SeedableRng::from_seed(seed);
        let mut sorter = RadixSorter::new();

        // pending edges and node updates.
        let mut edge_list = Vec::new();
        let mut node_lists = HashMap::new();

        // graph data; offsets into targets.
        let mut offsets = Vec::new();
        let mut targets = Vec::new();

        // holds the bfs parent of each node, or u32::max_value() if unset.
        let mut done = vec![u32::max_value(); 1 + (nodes / peers)];

        let start = time::precise_time_s();

        root.scoped(move |scope| {

            // generate part of a random graph.
            let graph = (0..edges / peers)
                .map(move |_| (rng.gen_range(0u32, nodes as u32), rng.gen_range(0u32, nodes as u32)))
                .to_stream(scope);

            // define a loop variable, for the (node, root) pairs.
            let (handle, stream) = scope.loop_variable(usize::max_value(), 1);

            // use the stream of edges
            graph.binary_notify(
                &stream,
                Exchange::new(|x: &(u32, u32)| x.0 as u64),
                Exchange::new(|x: &(u32, u32)| x.0 as u64),
                "BFS",
                vec![],
                move |input1, input2, output, notify| {

                    // receive edges, start to sort them
                    input1.for_each(|time, data| {
                        notify.notify_at(time);
                        edge_list.push(data.replace_with(Vec::new()));
                    });

                    // receive (node, root) pairs, note any new ones.
                    input2.for_each(|time, data| {
                        node_lists.entry(time.time())
                                  .or_insert_with(|| {
                                      notify.notify_at(time.clone());
                                      Vec::new()
                                  })
                                  .push(data.replace_with(Vec::new()));
                    });

                    notify.for_each(|time, _num, _notify| {

                        // maybe process the graph
                        if time.inner == 0 {

                            // print some diagnostic timing information
                            if index == 0 { println!("{}:\tsorting", time::precise_time_s() - start); }

                            // sort the edges
                            sorter.sort(&mut edge_list, &|x| x.0);

                            let mut count = 0;
                            for buffer in &edge_list { count += buffer.len(); }

                            // allocate sufficient memory, to avoid resizing.
                            offsets = Vec::with_capacity(1 + (nodes / peers));
                            targets = Vec::with_capacity(count);

                            // construct the graph
                            offsets.push(0);
                            let mut prev_node = 0;
                            for buffer in edge_list.drain(..) {
                                for (node, edge) in buffer {
                                    let temp = node / peers as u32;
                                    while prev_node < temp {
                                        prev_node += 1;
                                        offsets.push(targets.len() as u32)
                                    }
                                    targets.push(edge);
                                }
                            }
                            offsets.push(targets.len() as u32);
                        }

                        // print some diagnostic timing information
                        if index == 0 { println!("{}:\ttime: {:?}", time::precise_time_s() - start, time.time()); }

                        if let Some(mut todo) = node_lists.remove(&time) {
                            let mut session = output.session(&time);

                            // we could sort these, or not. try it out!
                            // sorter.sort(&mut todo, &|x| x.0);

                            for buffer in todo.drain(..) {
                                for (node, prev) in buffer {
                                    let temp = (node as usize) / peers;
                                    if done[temp] == u32::max_value() {
                                        done[temp] = prev;
                                        let lower = offsets[temp] as usize;
                                        let upper = offsets[temp + 1] as usize;
                                        for &target in &targets[lower..upper] {
                                            session.give((target, node));
                                        }
                                    }
                                }
                            }
                        }
                    });
                }
            )
            .concat(&(0..1).map(|x| (x,x)).to_stream(scope))
            .connect_loop(handle);
        });
    }).unwrap(); // asserts error-free execution;
}
<commit_msg>correct use of RadixSort<commit_after>extern crate time;
extern crate rand;
extern crate timely;
extern crate timely_sort;

use std::collections::HashMap;

use rand::{Rng, SeedableRng, StdRng};
use timely_sort::{RadixSorter, RadixSorterBase};
use timely_sort::LSBRadixSorter as Sorter;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::dataflow::channels::pact::Exchange;
fn main() {

    // command-line args: numbers of nodes and edges in the random graph.
    let nodes: usize = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();

    timely::execute_from_args(std::env::args().skip(3), move |root| {

        let index = root.index();
        let peers = root.peers();

        let seed: &[_] = &[1, 2, 3, root.index()];
        let mut rng: StdRng = SeedableRng::from_seed(seed);
        let mut sorter = Sorter::new();

        // pending edges and node updates.
        let mut edge_list = Vec::new();
        let mut node_lists = HashMap::new();

        // graph data; offsets into targets.
        let mut offsets = Vec::new();
        let mut targets = Vec::new();

        // holds the bfs parent of each node, or u32::max_value() if unset.
        let mut done = vec![u32::max_value(); 1 + (nodes / peers)];

        let start = time::precise_time_s();

        root.scoped(move |scope| {

            // generate part of a random graph.
            let graph = (0..edges / peers)
                .map(move |_| (rng.gen_range(0u32, nodes as u32), rng.gen_range(0u32, nodes as u32)))
                .to_stream(scope);

            // define a loop variable, for the (node, root) pairs.
            let (handle, stream) = scope.loop_variable(usize::max_value(), 1);

            // use the stream of edges
            graph.binary_notify(
                &stream,
                Exchange::new(|x: &(u32, u32)| x.0 as u64),
                Exchange::new(|x: &(u32, u32)| x.0 as u64),
                "BFS",
                vec![],
                move |input1, input2, output, notify| {

                    // receive edges, start to sort them
                    input1.for_each(|time, data| {
                        notify.notify_at(time);
                        edge_list.push(data.replace_with(Vec::new()));
                    });

                    // receive (node, root) pairs, note any new ones.
                    input2.for_each(|time, data| {
                        node_lists.entry(time.time())
                                  .or_insert_with(|| {
                                      notify.notify_at(time.clone());
                                      Vec::new()
                                  })
                                  .push(data.replace_with(Vec::new()));
                    });

                    notify.for_each(|time, _num, _notify| {

                        // maybe process the graph
                        if time.inner == 0 {

                            // print some diagnostic timing information
                            if index == 0 { println!("{}:\tsorting", time::precise_time_s() - start); }

                            // sort the edges
                            sorter.sort(&mut edge_list, &|x| x.0);

                            let mut count = 0;
                            for buffer in &edge_list { count += buffer.len(); }

                            // allocate sufficient memory, to avoid resizing.
                            offsets = Vec::with_capacity(1 + (nodes / peers));
                            targets = Vec::with_capacity(count);

                            // construct the graph
                            offsets.push(0);
                            let mut prev_node = 0;
                            for buffer in edge_list.drain(..) {
                                for (node, edge) in buffer {
                                    let temp = node / peers as u32;
                                    while prev_node < temp {
                                        prev_node += 1;
                                        offsets.push(targets.len() as u32)
                                    }
                                    targets.push(edge);
                                }
                            }
                            offsets.push(targets.len() as u32);
                        }

                        // print some diagnostic timing information
                        if index == 0 { println!("{}:\ttime: {:?}", time::precise_time_s() - start, time.time()); }

                        if let Some(mut todo) = node_lists.remove(&time) {
                            let mut session = output.session(&time);

                            // we could sort these, or not. try it out!
                            // sorter.sort(&mut todo, &|x| x.0);

                            for buffer in todo.drain(..) {
                                for (node, prev) in buffer {
                                    let temp = (node as usize) / peers;
                                    if done[temp] == u32::max_value() {
                                        done[temp] = prev;
                                        let lower = offsets[temp] as usize;
                                        let upper = offsets[temp + 1] as usize;
                                        for &target in &targets[lower..upper] {
                                            session.give((target, node));
                                        }
                                    }
                                }
                            }
                        }
                    });
                }
            )
            .concat(&(0..1).map(|x| (x,x)).to_stream(scope))
            .connect_loop(handle);
        });
    }).unwrap(); // asserts error-free execution;
}
<|endoftext|>


<commit_before>use super::{Region, RegionIndex};
use std::mem;
use rustc::infer::InferCtxt;
use rustc::mir::{Location, Mir};
use rustc_data_structures::indexed_vec::{Idx, IndexVec};
use rustc_data_structures::fx::FxHashSet;

pub struct InferenceContext {
    definitions: IndexVec<RegionIndex, VarDefinition>,
    constraints: IndexVec<ConstraintIndex, Constraint>,
    errors: IndexVec<InferenceErrorIndex, InferenceError>,
}

pub struct InferenceError {
    pub constraint_point: Location,
    pub name: (), // TODO(nashenas88) RegionName
}

newtype_index!(InferenceErrorIndex);

struct VarDefinition {
    name: (), // TODO(nashenas88) RegionName
    value: Region,
    capped: bool,
}

impl VarDefinition {
    pub fn new(value: Region) -> Self {
        Self {
            name: (),
            value,
            capped: false,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct Constraint {
    sub: RegionIndex,
    sup: RegionIndex,
    point: Location,
}

newtype_index!(ConstraintIndex);

impl InferenceContext {
    pub fn new(values: IndexVec<RegionIndex, Region>) -> Self {
        Self {
            definitions: values.into_iter().map(VarDefinition::new).collect(),
            constraints: IndexVec::new(),
            errors: IndexVec::new(),
        }
    }

    pub fn cap_var(&mut self, v: RegionIndex) {
        self.definitions[v].capped = true;
    }

    pub fn add_live_point(&mut self, v: RegionIndex, point: Location) {
        debug!("add_live_point({:?}, {:?})", v, point);
        let definition = &mut self.definitions[v];
        if definition.value.add_point(point) {
            if definition.capped {
                self.errors.push(InferenceError {
                    constraint_point: point,
                    name: definition.name,
                });
            }
        }
    }

    pub fn add_outlives(&mut self, sup: RegionIndex, sub: RegionIndex, point: Location) {
        debug!("add_outlives({:?}: {:?} @ {:?}", sup, sub, point);
        self.constraints.push(Constraint { sup, sub, point });
    }

    pub fn region(&self, v: RegionIndex) -> &Region {
        &self.definitions[v].value
    }

    pub fn solve<'a, 'gcx, 'tcx>(
        &mut self,
        infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,
        mir: &'a Mir<'tcx>,
    ) -> IndexVec<InferenceErrorIndex, InferenceError>
    where
        'gcx: 'tcx + 'a,
        'tcx: 'a,
    {
        let mut changed = true;
        let mut dfs = Dfs::new(infcx, mir);
        while changed {
            changed = false;
            for constraint in &self.constraints {
                let sub = &self.definitions[constraint.sub].value.clone();
                let sup_def = &mut self.definitions[constraint.sup];
                debug!("constraint: {:?}", constraint);
                debug!("    sub (before): {:?}", sub);
                debug!("    sup (before): {:?}", sup_def.value);

                if dfs.copy(sub, &mut sup_def.value, constraint.point) {
                    changed = true;
                    if sup_def.capped {
                        // This is kind of a hack, but when we add a
                        // constraint, the "point" is always the point
                        // AFTER the action that induced the
                        // constraint. So report the error on the
                        // action BEFORE that.
                        assert!(constraint.point.statement_index > 0);
                        let p = Location {
                            block: constraint.point.block,
                            statement_index: constraint.point.statement_index - 1,
                        };

                        self.errors.push(InferenceError {
                            constraint_point: p,
                            name: sup_def.name,
                        });
                    }
                }

                debug!("    sup (after) : {:?}", sup_def.value);
                debug!("    changed     : {:?}", changed);
            }
            debug!("\n");
        }

        mem::replace(&mut self.errors, IndexVec::new())
    }
}

struct Dfs<'a, 'gcx: 'tcx + 'a, 'tcx: 'a> {
    infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,
    mir: &'a Mir<'tcx>,
}

impl<'a, 'gcx: 'tcx, 'tcx: 'a> Dfs<'a, 'gcx, 'tcx> {
    fn new(infcx: &'a InferCtxt<'a, 'gcx, 'tcx>, mir: &'a Mir<'tcx>) -> Self {
        Self { infcx, mir }
    }

    fn copy(
        &mut self,
        from_region: &Region,
        to_region: &mut Region,
        start_point: Location,
    ) -> bool {
        let mut changed = false;

        let mut stack = vec![];
        let mut visited = FxHashSet();

        stack.push(start_point);
        while let Some(p) = stack.pop() {
            debug!("        dfs: p={:?}", p);

            if !from_region.may_contain(p) {
                debug!("            not in from-region");
                continue;
            }

            if !visited.insert(p) {
                debug!("            already visited");
                continue;
            }

            changed |= to_region.add_point(p);

            let block_data = &self.mir[p.block];
            let successor_points = if p.statement_index < block_data.statements.len() {
                vec![Location {
                    statement_index: p.statement_index + 1,
                    ..p
                }]
            } else {
                block_data.terminator()
                    .successors()
                    .iter()
                    .map(|&basic_block| Location {
                        statement_index: 0,
                        block: basic_block,
                    })
                    .collect::<Vec<_>>()
            };

            if successor_points.is_empty() {
                // TODO handle free regions
                // If we reach the END point in the graph, then copy
                // over any skolemized end points in the `from_region`
                // and make sure they are included in the `to_region`.
                // for region_decl in self.infcx.tcx.tables.borrow().free_region_map() {
                //     // TODO(nashenas88) figure out skolemized_end points
                //     let block = self.env.graph.skolemized_end(region_decl.name);
                //     let skolemized_end_point = Location {
                //         block,
                //         statement_index: 0,
                //     };
                //     changed |= to_region.add_point(skolemized_end_point);
                // }
            } else {
                stack.extend(successor_points);
            }
        }

        changed
    }
}
<commit_msg>Add allow_dead code directive to unused methods<commit_after>use super::{Region, RegionIndex};
use std::mem;
use rustc::infer::InferCtxt;
use rustc::mir::{Location, Mir};
use rustc_data_structures::indexed_vec::{Idx, IndexVec};
use rustc_data_structures::fx::FxHashSet;

pub struct InferenceContext {
    definitions: IndexVec<RegionIndex, VarDefinition>,
    constraints: IndexVec<ConstraintIndex, Constraint>,
    errors: IndexVec<InferenceErrorIndex, InferenceError>,
}

pub struct InferenceError {
    pub constraint_point: Location,
    pub name: (), // TODO(nashenas88) RegionName
}

newtype_index!(InferenceErrorIndex);

struct VarDefinition {
    name: (), // TODO(nashenas88) RegionName
    value: Region,
    capped: bool,
}

impl VarDefinition {
    pub fn new(value: Region) -> Self {
        Self {
            name: (),
            value,
            capped: false,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct Constraint {
    sub: RegionIndex,
    sup: RegionIndex,
    point: Location,
}

newtype_index!(ConstraintIndex);

impl InferenceContext {
    pub fn new(values: IndexVec<RegionIndex, Region>) -> Self {
        Self {
            definitions: values.into_iter().map(VarDefinition::new).collect(),
            constraints: IndexVec::new(),
            errors: IndexVec::new(),
        }
    }

    #[allow(dead_code)]
    pub fn cap_var(&mut self, v: RegionIndex) {
        self.definitions[v].capped = true;
    }

    #[allow(dead_code)]
    pub fn add_live_point(&mut self, v: RegionIndex, point: Location) {
        debug!("add_live_point({:?}, {:?})", v, point);
        let definition = &mut self.definitions[v];
        if definition.value.add_point(point) {
            if definition.capped {
                self.errors.push(InferenceError {
                    constraint_point: point,
                    name: definition.name,
                });
            }
        }
    }

    #[allow(dead_code)]
    pub fn add_outlives(&mut self, sup: RegionIndex, sub: RegionIndex, point: Location) {
        debug!("add_outlives({:?}: {:?} @ {:?}", sup, sub, point);
        self.constraints.push(Constraint { sup, sub, point });
    }

    #[allow(dead_code)]
    pub fn region(&self, v: RegionIndex) -> &Region {
        &self.definitions[v].value
    }

    pub fn solve<'a, 'gcx, 'tcx>(
        &mut self,
        infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,
        mir: &'a Mir<'tcx>,
    ) -> IndexVec<InferenceErrorIndex, InferenceError>
    where
        'gcx: 'tcx + 'a,
        'tcx: 'a,
    {
        let mut changed = true;
        let mut dfs = Dfs::new(infcx, mir);
        while changed {
            changed = false;
            for constraint in &self.constraints {
                let sub = &self.definitions[constraint.sub].value.clone();
                let sup_def = &mut self.definitions[constraint.sup];
                debug!("constraint: {:?}", constraint);
                debug!("    sub (before): {:?}", sub);
                debug!("    sup (before): {:?}", sup_def.value);

                if dfs.copy(sub, &mut sup_def.value, constraint.point) {
                    changed = true;
                    if sup_def.capped {
                        // This is kind of a hack, but when we add a
                        // constraint, the "point" is always the point
                        // AFTER the action that induced the
                        // constraint. So report the error on the
                        // action BEFORE that.
                        assert!(constraint.point.statement_index > 0);
                        let p = Location {
                            block: constraint.point.block,
                            statement_index: constraint.point.statement_index - 1,
                        };

                        self.errors.push(InferenceError {
                            constraint_point: p,
                            name: sup_def.name,
                        });
                    }
                }

                debug!("    sup (after) : {:?}", sup_def.value);
                debug!("    changed     : {:?}", changed);
            }
            debug!("\n");
        }

        mem::replace(&mut self.errors, IndexVec::new())
    }
}

struct Dfs<'a, 'gcx: 'tcx + 'a, 'tcx: 'a> {
    #[allow(dead_code)]
    infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,
    mir: &'a Mir<'tcx>,
}

impl<'a, 'gcx: 'tcx, 'tcx: 'a> Dfs<'a, 'gcx, 'tcx> {
    fn new(infcx: &'a InferCtxt<'a, 'gcx, 'tcx>, mir: &'a Mir<'tcx>) -> Self {
        Self { infcx, mir }
    }

    fn copy(
        &mut self,
        from_region: &Region,
        to_region: &mut Region,
        start_point: Location,
    ) -> bool {
        let mut changed = false;

        let mut stack = vec![];
        let mut visited = FxHashSet();

        stack.push(start_point);
        while let Some(p) = stack.pop() {
            debug!("        dfs: p={:?}", p);

            if !from_region.may_contain(p) {
                debug!("            not in from-region");
                continue;
            }

            if !visited.insert(p) {
                debug!("            already visited");
                continue;
            }

            changed |= to_region.add_point(p);

            let block_data = &self.mir[p.block];
            let successor_points = if p.statement_index < block_data.statements.len() {
                vec![Location {
                    statement_index: p.statement_index + 1,
                    ..p
                }]
            } else {
                block_data.terminator()
                    .successors()
                    .iter()
                    .map(|&basic_block| Location {
                        statement_index: 0,
                        block: basic_block,
                    })
                    .collect::<Vec<_>>()
            };

            if successor_points.is_empty() {
                // TODO handle free regions
                // If we reach the END point in the graph, then copy
                // over any skolemized end points in the `from_region`
                // and make sure they are included in the `to_region`.
                // for region_decl in self.infcx.tcx.tables.borrow().free_region_map() {
                //     // TODO(nashenas88) figure out skolemized_end points
                //     let block = self.env.graph.skolemized_end(region_decl.name);
                //     let skolemized_end_point = Location {
                //         block,
                //         statement_index: 0,
                //     };
                //     changed |= to_region.add_point(skolemized_end_point);
                // }
            } else {
                stack.extend(successor_points);
            }
        }

        changed
    }
}
<|endoftext|>


<commit_before>use credentials::Key;
use base64;
use ring::constant_time;
use std::io::Write;
use std::ops::Deref;
use error::*;
use time;

/// The kind of MAC calcuation (corresponding to the first line of the message)
pub enum MacType {
    Header,
    Response,
    Bewit,
}

/// Mac represents a message authentication code, the signature in a Hawk transaction.
///
/// This class supports creating Macs using the Hawk specification, and comparing Macs
/// using a cosntant-time comparison (thus preventing timing side-channel attacks).
#[derive(Debug, Clone)]
pub struct Mac(Vec<u8>);

impl Mac {
    pub fn new(mac_type: MacType,
               key: &Key,
               ts: time::Timespec,
               nonce: &str,
               method: &str,
               host: &str,
               port: u16,
               path: &str,
               hash: Option<&[u8]>,
               ext: Option<&str>)
               -> Result<Mac> {
        let mut buffer: Vec<u8> = vec![];

        write!(buffer,
               "{}\n",
               match mac_type {
                   MacType::Header => "hawk.1.header",
                   MacType::Response => "hawk.1.response",
                   MacType::Bewit => "hawk.1.bewit",
               })?;
        write!(buffer, "{}\n", ts.sec)?;
        write!(buffer, "{}\n", nonce)?;
        write!(buffer, "{}\n", method)?;
        write!(buffer, "{}\n", path)?;
        write!(buffer, "{}\n", host)?;
        write!(buffer, "{}\n", port)?;

        if let Some(h) = hash {
            write!(buffer, "{}\n", base64::encode(h))?;
        } else {
            write!(buffer, "\n")?;
        }

        match ext {
            Some(e) => write!(buffer, "{}\n", e)?,
            None => write!(buffer, "\n")?,
        };

        println!("{:?}", String::from_utf8(buffer.clone()));

        Ok(Mac(key.sign(buffer.as_ref())))
    }
}

impl AsRef<[u8]> for Mac {
    fn as_ref(&self) -> &[u8] {
        &self.0[..]
    }
}

impl From<Vec<u8>> for Mac {
    fn from(original: Vec<u8>) -> Self {
        Mac(original)
    }
}

impl Deref for Mac {
    type Target = Vec<u8>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl PartialEq for Mac {
    fn eq(&self, other: &Mac) -> bool {
        match constant_time::verify_slices_are_equal(&self.0[..], &other.0[..]) {
            Ok(_) => true,
            Err(_) => false,
        }
    }
}

#[cfg(test)]
mod test {
    use super::{Mac, MacType};
    use time::Timespec;
    use credentials::Key;
    use ring::digest;

    fn key() -> Key {
        Key::new(vec![11u8, 19, 228, 209, 79, 189, 200, 59, 166, 47, 86, 254, 235, 184, 120, 197,
                      75, 152, 201, 79, 115, 61, 111, 242, 219, 187, 173, 14, 227, 108, 60, 232],
                 &digest::SHA256)
    }

    #[test]
    fn test_make_mac() {
        let key = key();
        let mac = Mac::new(MacType::Header,
                           &key,
                           Timespec::new(1000, 100),
                           "nonny",
                           "POST",
                           "mysite.com",
                           443,
                           "/v1/api",
                           None,
                           None)
            .unwrap();
        println!("got {:?}", mac);
        assert!(mac.0 ==
                vec![192, 227, 235, 121, 157, 185, 197, 79, 189, 214, 235, 139, 9, 232, 99, 55,
                     67, 30, 68, 0, 150, 187, 192, 238, 21, 200, 209, 107, 245, 159, 243, 178]);
    }

    #[test]
    fn test_make_mac_hash() {
        let key = key();
        let hash = vec![1, 2, 3, 4, 5];
        let mac = Mac::new(MacType::Header,
                           &key,
                           Timespec::new(1000, 100),
                           "nonny",
                           "POST",
                           "mysite.com",
                           443,
                           "/v1/api",
                           Some(&hash),
                           None)
            .unwrap();
        println!("got {:?}", mac);
        assert!(mac.0 ==
                vec![61, 128, 208, 253, 88, 135, 190, 196, 1, 69, 153, 193, 124, 4, 195, 87, 38,
                     96, 181, 34, 65, 234, 58, 157, 175, 175, 145, 151, 61, 0, 57, 5]);
    }

    #[test]
    fn test_make_mac_ext() {
        let key = key();
        let ext = "ext-data".to_string();
        let mac = Mac::new(MacType::Header,
                           &key,
                           Timespec::new(1000, 100),
                           "nonny",
                           "POST",
                           "mysite.com",
                           443,
                           "/v1/api",
                           None,
                           Some(&ext))
            .unwrap();
        println!("got {:?}", mac);
        assert!(mac.0 ==
                vec![187, 104, 238, 100, 168, 112, 37, 68, 187, 141, 168, 155, 177, 193, 113, 0,
                     50, 105, 127, 36, 24, 117, 200, 251, 138, 199, 108, 14, 105, 123, 234, 119]);
    }
}
<commit_msg>remove a println! in production code<commit_after>use credentials::Key;
use base64;
use ring::constant_time;
use std::io::Write;
use std::ops::Deref;
use error::*;
use time;

/// The kind of MAC calcuation (corresponding to the first line of the message)
pub enum MacType {
    Header,
    Response,
    Bewit,
}

/// Mac represents a message authentication code, the signature in a Hawk transaction.
///
/// This class supports creating Macs using the Hawk specification, and comparing Macs
/// using a cosntant-time comparison (thus preventing timing side-channel attacks).
#[derive(Debug, Clone)]
pub struct Mac(Vec<u8>);

impl Mac {
    pub fn new(mac_type: MacType,
               key: &Key,
               ts: time::Timespec,
               nonce: &str,
               method: &str,
               host: &str,
               port: u16,
               path: &str,
               hash: Option<&[u8]>,
               ext: Option<&str>)
               -> Result<Mac> {
        let mut buffer: Vec<u8> = vec![];

        write!(buffer,
               "{}\n",
               match mac_type {
                   MacType::Header => "hawk.1.header",
                   MacType::Response => "hawk.1.response",
                   MacType::Bewit => "hawk.1.bewit",
               })?;
        write!(buffer, "{}\n", ts.sec)?;
        write!(buffer, "{}\n", nonce)?;
        write!(buffer, "{}\n", method)?;
        write!(buffer, "{}\n", path)?;
        write!(buffer, "{}\n", host)?;
        write!(buffer, "{}\n", port)?;

        if let Some(h) = hash {
            write!(buffer, "{}\n", base64::encode(h))?;
        } else {
            write!(buffer, "\n")?;
        }

        match ext {
            Some(e) => write!(buffer, "{}\n", e)?,
            None => write!(buffer, "\n")?,
        };

        Ok(Mac(key.sign(buffer.as_ref())))
    }
}

impl AsRef<[u8]> for Mac {
    fn as_ref(&self) -> &[u8] {
        &self.0[..]
    }
}

impl From<Vec<u8>> for Mac {
    fn from(original: Vec<u8>) -> Self {
        Mac(original)
    }
}

impl Deref for Mac {
    type Target = Vec<u8>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl PartialEq for Mac {
    fn eq(&self, other: &Mac) -> bool {
        match constant_time::verify_slices_are_equal(&self.0[..], &other.0[..]) {
            Ok(_) => true,
            Err(_) => false,
        }
    }
}

#[cfg(test)]
mod test {
    use super::{Mac, MacType};
    use time::Timespec;
    use credentials::Key;
    use ring::digest;

    fn key() -> Key {
        Key::new(vec![11u8, 19, 228, 209, 79, 189, 200, 59, 166, 47, 86, 254, 235, 184, 120, 197,
                      75, 152, 201, 79, 115, 61, 111, 242, 219, 187, 173, 14, 227, 108, 60, 232],
                 &digest::SHA256)
    }

    #[test]
    fn test_make_mac() {
        let key = key();
        let mac = Mac::new(MacType::Header,
                           &key,
                           Timespec::new(1000, 100),
                           "nonny",
                           "POST",
                           "mysite.com",
                           443,
                           "/v1/api",
                           None,
                           None)
            .unwrap();
        println!("got {:?}", mac);
        assert!(mac.0 ==
                vec![192, 227, 235, 121, 157, 185, 197, 79, 189, 214, 235, 139, 9, 232, 99, 55,
                     67, 30, 68, 0, 150, 187, 192, 238, 21, 200, 209, 107, 245, 159, 243, 178]);
    }

    #[test]
    fn test_make_mac_hash() {
        let key = key();
        let hash = vec![1, 2, 3, 4, 5];
        let mac = Mac::new(MacType::Header,
                           &key,
                           Timespec::new(1000, 100),
                           "nonny",
                           "POST",
                           "mysite.com",
                           443,
                           "/v1/api",
                           Some(&hash),
                           None)
            .unwrap();
        println!("got {:?}", mac);
        assert!(mac.0 ==
                vec![61, 128, 208, 253, 88, 135, 190, 196, 1, 69, 153, 193, 124, 4, 195, 87, 38,
                     96, 181, 34, 65, 234, 58, 157, 175, 175, 145, 151, 61, 0, 57, 5]);
    }

    #[test]
    fn test_make_mac_ext() {
        let key = key();
        let ext = "ext-data".to_string();
        let mac = Mac::new(MacType::Header,
                           &key,
                           Timespec::new(1000, 100),
                           "nonny",
                           "POST",
                           "mysite.com",
                           443,
                           "/v1/api",
                           None,
                           Some(&ext))
            .unwrap();
        println!("got {:?}", mac);
        assert!(mac.0 ==
                vec![187, 104, 238, 100, 168, 112, 37, 68, 187, 141, 168, 155, 177, 193, 113, 0,
                     50, 105, 127, 36, 24, 117, 200, 251, 138, 199, 108, 14, 105, 123, 234, 119]);
    }
}
<|endoftext|>


<commit_before>#![allow(
    clippy::blocks_in_if_conditions,
    clippy::cast_lossless,
    clippy::cast_possible_truncation,
    clippy::manual_map,
    clippy::map_unwrap_or,
    clippy::needless_pass_by_value,
    clippy::option_if_let_else,
    clippy::range_plus_one,
    clippy::single_match_else,
    clippy::too_many_lines,
    clippy::wrong_self_convention
)]

extern crate proc_macro;

mod ast;
mod attr;
mod expand;
mod fmt;
mod generics;
mod prop;
mod valid;

use proc_macro::TokenStream;
use syn::{parse_macro_input, DeriveInput};

#[proc_macro_derive(Error, attributes(backtrace, error, from, source))]
pub fn derive_error(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    expand::derive(&input)
        .unwrap_or_else(|err| err.to_compile_error())
        .into()
}
<commit_msg>Ignore manual_find clippy lint<commit_after>#![allow(
    clippy::blocks_in_if_conditions,
    clippy::cast_lossless,
    clippy::cast_possible_truncation,
    clippy::manual_find,
    clippy::manual_map,
    clippy::map_unwrap_or,
    clippy::needless_pass_by_value,
    clippy::option_if_let_else,
    clippy::range_plus_one,
    clippy::single_match_else,
    clippy::too_many_lines,
    clippy::wrong_self_convention
)]

extern crate proc_macro;

mod ast;
mod attr;
mod expand;
mod fmt;
mod generics;
mod prop;
mod valid;

use proc_macro::TokenStream;
use syn::{parse_macro_input, DeriveInput};

#[proc_macro_derive(Error, attributes(backtrace, error, from, source))]
pub fn derive_error(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    expand::derive(&input)
        .unwrap_or_else(|err| err.to_compile_error())
        .into()
}
<|endoftext|>


<commit_before>use crate::utils::span_lint_and_then;
use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::{Expr, ExprKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_session::{declare_lint_pass, declare_tool_lint};
use rustc_span::source_map::Spanned;

use crate::utils::{match_def_path, paths};
use rustc_ast::ast::LitKind;
use rustc_hir as hir;

declare_clippy_lint! {
    /// **What it does:** Finds occurences of `Vec::resize(0, an_int)`
    ///
    /// **Why is this bad?** This is probably an argument inversion mistake.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// vec!(1, 2, 3, 4, 5).resize(0, 5)
    /// ```
    pub VEC_RESIZE_TO_ZERO,
    correctness,
    "emptying a vector with `resize(0, an_int)` instead of `clear()` is probably an argument inversion mistake"
}

declare_lint_pass!(VecResizeToZero => [VEC_RESIZE_TO_ZERO]);

impl<'tcx> LateLintPass<'tcx> for VecResizeToZero {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
        if_chain! {
            if let hir::ExprKind::MethodCall(path_segment, _, ref args, _) = expr.kind;
            if let Some(method_def_id) = cx.tables().type_dependent_def_id(expr.hir_id);
            if match_def_path(cx, method_def_id, &paths::VEC_RESIZE) && args.len() == 3;
            if let ExprKind::Lit(Spanned { node: LitKind::Int(0, _), .. }) = args[1].kind;
            if let ExprKind::Lit(Spanned { node: LitKind::Int(..), .. }) = args[2].kind;
            then {
                let method_call_span = expr.span.with_lo(path_segment.ident.span.lo());
                span_lint_and_then(
                    cx,
                    VEC_RESIZE_TO_ZERO,
                    expr.span,
                    "emptying a vector with `resize`",
                    |db| {
                        db.help("the arguments may be inverted...");
                        db.span_suggestion(
                            method_call_span,
                            "...or you can empty the vector with",
                            "clear()".to_string(),
                            Applicability::MaybeIncorrect,
                        );
                    },
                );
            }
        }
    }
}
<commit_msg>Auto merge of #5814 - giraffate:fix_typo, r=flip1995<commit_after>use crate::utils::span_lint_and_then;
use if_chain::if_chain;
use rustc_errors::Applicability;
use rustc_hir::{Expr, ExprKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_session::{declare_lint_pass, declare_tool_lint};
use rustc_span::source_map::Spanned;

use crate::utils::{match_def_path, paths};
use rustc_ast::ast::LitKind;
use rustc_hir as hir;

declare_clippy_lint! {
    /// **What it does:** Finds occurrences of `Vec::resize(0, an_int)`
    ///
    /// **Why is this bad?** This is probably an argument inversion mistake.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// vec!(1, 2, 3, 4, 5).resize(0, 5)
    /// ```
    pub VEC_RESIZE_TO_ZERO,
    correctness,
    "emptying a vector with `resize(0, an_int)` instead of `clear()` is probably an argument inversion mistake"
}

declare_lint_pass!(VecResizeToZero => [VEC_RESIZE_TO_ZERO]);

impl<'tcx> LateLintPass<'tcx> for VecResizeToZero {
    fn check_expr(&mut self, cx: &LateContext<'tcx>, expr: &'tcx Expr<'_>) {
        if_chain! {
            if let hir::ExprKind::MethodCall(path_segment, _, ref args, _) = expr.kind;
            if let Some(method_def_id) = cx.tables().type_dependent_def_id(expr.hir_id);
            if match_def_path(cx, method_def_id, &paths::VEC_RESIZE) && args.len() == 3;
            if let ExprKind::Lit(Spanned { node: LitKind::Int(0, _), .. }) = args[1].kind;
            if let ExprKind::Lit(Spanned { node: LitKind::Int(..), .. }) = args[2].kind;
            then {
                let method_call_span = expr.span.with_lo(path_segment.ident.span.lo());
                span_lint_and_then(
                    cx,
                    VEC_RESIZE_TO_ZERO,
                    expr.span,
                    "emptying a vector with `resize`",
                    |db| {
                        db.help("the arguments may be inverted...");
                        db.span_suggestion(
                            method_call_span,
                            "...or you can empty the vector with",
                            "clear()".to_string(),
                            Applicability::MaybeIncorrect,
                        );
                    },
                );
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Utilities for generating and reading the self-signed certificate for use with the control
//! gateway.

use crate::{crypto::keys::NamedRevision,
            tls::rustls_wrapper::{self,
                                  Error as RustlsReadersError}};
use rcgen::{Certificate as RcgenCertificate,
            CertificateParams,
            DistinguishedName,
            DnType,
            RcgenError,
            PKCS_ECDSA_P256_SHA256};
use rustls::{Certificate,
             PrivateKey,
             RootCertStore};
use std::{fs::{self,
               File},
          io::{Error as IoError,
               Write},
          path::{Path,
                 PathBuf}};
use thiserror::Error;

const NAME_PREFIX: &str = "ctl-gateway";
const CRT_EXTENSION: &str = "crt.pem";
const KEY_EXTENSION: &str = "key.pem";

#[derive(Error, Debug)]
pub enum Error {
    #[error("ctl gateway certificate lookup failed when trying to match files {0}")]
    FailedToMatchPattern(String),
    #[error("ctl gateway certificate lookup failed, err: {0}")]
    RustlsReaders(#[from] RustlsReadersError),
    #[error("ctl gateway certificate generation failed, err: {0}")]
    CertificateGeneration(#[from] RcgenError),
    #[error("writing the ctl gateway certificate failed, err: {0}")]
    CertificateWrite(#[from] IoError),
}

pub fn generate_self_signed_certificate_and_key(subject_alternate_name: &str,
                                                path: impl AsRef<Path>)
                                                -> Result<(), Error> {
    let mut params = CertificateParams::new(vec![subject_alternate_name.to_string(),
                                                 "localhost".to_string(),
                                                 "127.0.0.1".to_string(),
                                                 "0.0.0.0".to_string()]);
    let mut distinguished_name = DistinguishedName::new();
    distinguished_name.push(DnType::CommonName, "Habitat Supervisor Control Gateway");
    params.distinguished_name = distinguished_name;
    params.alg = &PKCS_ECDSA_P256_SHA256;

    let certificate = RcgenCertificate::from_params(params)?;
    let crt = certificate.serialize_pem()?;
    let key = certificate.serialize_private_key_pem();

    fs::create_dir_all(&path)?;
    let named_revision = NamedRevision::new(NAME_PREFIX.to_string());

    let crt_path = path.as_ref()
                       .join(format!("{}.{}", named_revision, CRT_EXTENSION));
    let mut crt_file = File::create(crt_path)?;
    crt_file.write_all(crt.as_bytes())?;

    let key_path = path.as_ref()
                       .join(format!("{}.{}", named_revision, KEY_EXTENSION));
    let mut key_file = File::create(key_path)?;
    key_file.write_all(key.as_bytes())?;

    Ok(())
}

fn get_latest_path(path: impl AsRef<Path>, pattern: &str) -> Result<PathBuf, Error> {
    let pattern = path.as_ref().join(pattern);
    let pattern = pattern.to_string_lossy();
    glob::glob(&pattern).expect("valid pattern")
                        .filter_map(std::result::Result::ok)
                        .filter(|p| p.metadata().map(|m| m.is_file()).unwrap_or(false))
                        .max()
                        .ok_or_else(|| Error::FailedToMatchPattern(pattern.to_string()))
}

pub fn latest_certificates(path: impl AsRef<Path>) -> Result<Vec<Certificate>, Error> {
    let path = get_latest_path(path, &format!("{}-*.{}", NAME_PREFIX, CRT_EXTENSION))?;
    Ok(rustls_wrapper::certificates_from_file(&path)?)
}

pub fn latest_private_key(path: impl AsRef<Path>) -> Result<PrivateKey, Error> {
    let path = get_latest_path(path, &format!("{}-*.{}", NAME_PREFIX, KEY_EXTENSION))?;
    Ok(rustls_wrapper::private_key_from_file(&path)?)
}

pub fn latest_root_certificate_store(path: impl AsRef<Path>) -> Result<RootCertStore, Error> {
    let path = get_latest_path(path, &format!("{}-*.{}", NAME_PREFIX, CRT_EXTENSION))?;
    Ok(rustls_wrapper::root_certificate_store_from_file(&path)?)
}
<commit_msg>Remove ip addresses from ctl gateway certs<commit_after>//! Utilities for generating and reading the self-signed certificate for use with the control
//! gateway.

use crate::{crypto::keys::NamedRevision,
            tls::rustls_wrapper::{self,
                                  Error as RustlsReadersError}};
use rcgen::{Certificate as RcgenCertificate,
            CertificateParams,
            DistinguishedName,
            DnType,
            RcgenError,
            PKCS_ECDSA_P256_SHA256};
use rustls::{Certificate,
             PrivateKey,
             RootCertStore};
use std::{fs::{self,
               File},
          io::{Error as IoError,
               Write},
          path::{Path,
                 PathBuf}};
use thiserror::Error;

const NAME_PREFIX: &str = "ctl-gateway";
const CRT_EXTENSION: &str = "crt.pem";
const KEY_EXTENSION: &str = "key.pem";

#[derive(Error, Debug)]
pub enum Error {
    #[error("ctl gateway certificate lookup failed when trying to match files {0}")]
    FailedToMatchPattern(String),
    #[error("ctl gateway certificate lookup failed, err: {0}")]
    RustlsReaders(#[from] RustlsReadersError),
    #[error("ctl gateway certificate generation failed, err: {0}")]
    CertificateGeneration(#[from] RcgenError),
    #[error("writing the ctl gateway certificate failed, err: {0}")]
    CertificateWrite(#[from] IoError),
}

pub fn generate_self_signed_certificate_and_key(subject_alternate_name: &str,
                                                path: impl AsRef<Path>)
                                                -> Result<(), Error> {
    let mut params =
        CertificateParams::new(vec![subject_alternate_name.to_string(), "localhost".to_string(),]);
    let mut distinguished_name = DistinguishedName::new();
    distinguished_name.push(DnType::CommonName, "Habitat Supervisor Control Gateway");
    params.distinguished_name = distinguished_name;
    params.alg = &PKCS_ECDSA_P256_SHA256;

    let certificate = RcgenCertificate::from_params(params)?;
    let crt = certificate.serialize_pem()?;
    let key = certificate.serialize_private_key_pem();

    fs::create_dir_all(&path)?;
    let named_revision = NamedRevision::new(NAME_PREFIX.to_string());

    let crt_path = path.as_ref()
                       .join(format!("{}.{}", named_revision, CRT_EXTENSION));
    let mut crt_file = File::create(crt_path)?;
    crt_file.write_all(crt.as_bytes())?;

    let key_path = path.as_ref()
                       .join(format!("{}.{}", named_revision, KEY_EXTENSION));
    let mut key_file = File::create(key_path)?;
    key_file.write_all(key.as_bytes())?;

    Ok(())
}

fn get_latest_path(path: impl AsRef<Path>, pattern: &str) -> Result<PathBuf, Error> {
    let pattern = path.as_ref().join(pattern);
    let pattern = pattern.to_string_lossy();
    glob::glob(&pattern).expect("valid pattern")
                        .filter_map(std::result::Result::ok)
                        .filter(|p| p.metadata().map(|m| m.is_file()).unwrap_or(false))
                        .max()
                        .ok_or_else(|| Error::FailedToMatchPattern(pattern.to_string()))
}

pub fn latest_certificates(path: impl AsRef<Path>) -> Result<Vec<Certificate>, Error> {
    let path = get_latest_path(path, &format!("{}-*.{}", NAME_PREFIX, CRT_EXTENSION))?;
    Ok(rustls_wrapper::certificates_from_file(&path)?)
}

pub fn latest_private_key(path: impl AsRef<Path>) -> Result<PrivateKey, Error> {
    let path = get_latest_path(path, &format!("{}-*.{}", NAME_PREFIX, KEY_EXTENSION))?;
    Ok(rustls_wrapper::private_key_from_file(&path)?)
}

pub fn latest_root_certificate_store(path: impl AsRef<Path>) -> Result<RootCertStore, Error> {
    let path = get_latest_path(path, &format!("{}-*.{}", NAME_PREFIX, CRT_EXTENSION))?;
    Ok(rustls_wrapper::root_certificate_store_from_file(&path)?)
}
<|endoftext|>


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

//! Encapsulate running the `hab-sup` executable for tests.

use std::collections::HashSet;
use std::env;
use std::net::TcpListener;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::string::ToString;
use std::sync::Mutex;
use std::thread;
use std::time::Duration;

use hcore::url::BLDR_URL_ENVVAR;
use rand;
use rand::distributions::{IndependentSample, Range};

use super::test_butterfly;

lazy_static! {
    /// Keep track of all TCP ports currently being used by TestSup
    /// instances. Allows us to run tests in parallel without fear of
    /// port conflicts between them.
    static ref CLAIMED_PORTS: Mutex<HashSet<u16>> = {
        Mutex::new(HashSet::new())
    };
}

pub struct TestSup {
    pub hab_root: PathBuf,
    pub origin_name: String,
    pub package_name: String,
    pub service_group: String,
    pub http_port: u16,
    pub butterfly_port: u16,
    pub control_port: u16,
    pub butterfly_client: test_butterfly::Client,
    pub cmd: Command,
    pub process: Option<Child>,
}

/// Return a free TCP port number. We test to see that the system has
/// not already bound the port, while also tracking which ports are
/// being used by other test supervisors that may be running alongside
/// this one.
///
/// Once you receive a port number from this function, you can be
/// reasonably sure that you're the only one that will be using
/// it. There could be a race condition if the machine the tests are
/// running on just happens to claim the same port number for
/// something between the time we check and the time the TestSup
/// claims it. If that happens to you, you should probably buy lottery
/// tickets, though.
///
/// This function will recursively call itself with a decremented
/// value for `tries` if it happens to pick a port that's already in
/// use. Once all tries are used up, it panics! Yay!
fn unclaimed_port(tries: u16) -> u16 {
    if tries == 0 {
        panic!("Couldn't find an unclaimed port for the test Supervisor!")
    }
    let p = random_port();
    match TcpListener::bind(format!("127.0.0.1:{}", p)) {
        Ok(_listener) => {
            // The system hasn't bound it. Now we make sure none of
            // our other tests have bound it.
            let mut ports = CLAIMED_PORTS.lock().unwrap();
            if ports.contains(&p) {
                // Oops, another test is using it, try again
                thread::sleep(Duration::from_millis(500));
                unclaimed_port(tries - 1)
            } else {
                // Nobody was using it. Return the port; the TcpListener
                // that is currently bound to the port will be dropped,
                // thus freeing the port for our use.
                ports.insert(p);
                p
            }
        }
        Err(_) => {
            // port already in use, try again
            unclaimed_port(tries - 1)
        }
    }
}

/// Return a random unprivileged, unregistered TCP port number.
fn random_port() -> u16 {
    // IANA port registrations go to 49151
    let between = Range::new(49152, ::std::u16::MAX);
    let mut rng = rand::thread_rng();
    between.ind_sample(&mut rng)
}

/// Find an executable relative to the current integration testing
/// executable.
///
/// Thus if the current executable is
///
///    /home/me/habitat/target/debug/deps/compilation-ccaf2f45c24e3840
///
/// and we look for `hab-sup`, we'll find it at
///
///    /home/me/habitat/target/debug/hab-sup
///
fn find_exe<B>(binary_name: B) -> PathBuf
where
    B: AsRef<Path>,
{
    let exe_root = env::current_exe()
        .unwrap()
        .parent()  // deps
        .unwrap()
        .parent()  // debug
        .unwrap()
        .to_path_buf();
    let bin = exe_root.join(binary_name.as_ref());
    assert!(
        bin.exists(),
        format!(
            "Expected to find a {:?} executable at {:?}",
            binary_name.as_ref(),
            bin
        )
    );
    bin
}

/// Return whether or not the tests are being run with the `--nocapture` flag meaning we want to
/// see more output.
fn nocapture_set() -> bool {
    if env::args().any(|arg| arg == "--nocapture") {
        return true;
    } else {
        match env::var("RUST_TEST_NOCAPTURE") {
            Ok(val) => &val != "0",
            Err(_) => false,
        }
    }
}

impl TestSup {
    /// Create a new `TestSup` that will listen on randomly-selected
    /// ports for both gossip and HTTP requests so tests run in
    /// parallel don't step on each other.
    ///
    /// See also `new`.
    pub fn new_with_random_ports<R, O, P, S>(
        fs_root: R,
        origin: O,
        pkg_name: P,
        service_group: S,
    ) -> TestSup
    where
        R: AsRef<Path>,
        O: ToString,
        P: ToString,
        S: ToString,
    {
        // We'll give 10 tries to find a free port number
        let http_port = unclaimed_port(10);
        let butterfly_port = unclaimed_port(10);
        let control_port = unclaimed_port(10);

        TestSup::new(
            fs_root,
            origin,
            pkg_name,
            service_group,
            http_port,
            butterfly_port,
            control_port,
        )
    }

    /// Bundle up a Habitat Supervisor process along with an
    /// associated Butterfly client for injecting new configuration
    /// values. The Supervisor executable is the one that has been
    /// compiled for the current `cargo test` invocation.
    ///
    /// The Supervisor is configured to run a single package for a
    /// test. This package is assumed to have already been installed
    /// relative to `fs_root` (i.e., the `FS_ROOT` environment
    /// variable, which in our tests will be a randomly-named
    /// temporary directory that this Supervisor will view as `/`.).
    ///
    /// A Butterfly client is also created for interacting with this
    /// Supervisor and package. It is properly configured according to
    /// the value provided for `butterfly_port`. To use it, see the
    /// `apply_config` function.
    ///
    /// (No HTTP interaction with the Supervisor is currently called
    /// for, so we don't have a HTTP client.)
    pub fn new<R, O, P, S>(
        fs_root: R,
        origin: O,
        pkg_name: P,
        service_group: S,
        http_port: u16,
        butterfly_port: u16,
        control_port: u16,
    ) -> TestSup
    where
        R: AsRef<Path>,
        O: ToString,
        P: ToString,
        S: ToString,
    {
        let sup_exe = find_exe("hab-sup");
        let launcher_exe = find_exe("hab-launch");

        let mut cmd = Command::new(&launcher_exe);
        let listen_host = "0.0.0.0";

        let origin = origin.to_string();
        let pkg_name = pkg_name.to_string();
        let service_group = service_group.to_string();

        cmd.env(
            "TESTING_FS_ROOT",
            fs_root.as_ref().to_string_lossy().as_ref(),
        ).env("HAB_SUP_BINARY", &sup_exe)
            .env(BLDR_URL_ENVVAR, "http://hab.sup.test")
            .arg("run")
            .arg("--listen-gossip")
            .arg(format!("{}:{}", listen_host, butterfly_port))
            .arg("--listen-http")
            .arg(format!("{}:{}", listen_host, http_port))
            .arg("--listen-ctl")
            .arg(format!("{}:{}", listen_host, control_port))
            // Note: we will have already dropped off the spec files
            // needed to run our test service, so we don't supply a
            // package identifier here
            .stdin(Stdio::null());
        if !nocapture_set() {
            cmd.stdout(Stdio::null());
            cmd.stderr(Stdio::null());
        }

        let bc = test_butterfly::Client::new(&pkg_name, &service_group, butterfly_port);

        TestSup {
            hab_root: fs_root.as_ref().to_path_buf(),
            origin_name: origin,
            package_name: pkg_name,
            service_group: service_group.to_string(),
            http_port: http_port,
            butterfly_port: butterfly_port,
            control_port: control_port,
            butterfly_client: bc,
            cmd: cmd,
            process: None,
        }
    }

    /// Spawn a process actually running the Supervisor.
    pub fn start(&mut self) {
        let child = self.cmd.spawn().expect("Couldn't start the Supervisor!");
        self.process = Some(child);
    }

    /// The equivalent of performing `hab apply` with the given
    /// configuration.
    pub fn apply_config<T>(&mut self, toml_config: T)
    where
        T: ToString,
    {
        self.butterfly_client.apply(toml_config.to_string())
    }
}

// We kill the Supervisor so you don't have to! We also free up the
// ports used by this Supervisor so other tests can use them.
impl Drop for TestSup {
    fn drop(&mut self) {
        let mut ports = CLAIMED_PORTS.lock().unwrap();
        ports.remove(&self.http_port);
        ports.remove(&self.butterfly_port);
        ports.remove(&self.control_port);
        self.process
            .take()
            .expect("No process to kill!")
            .kill()
            .expect("Tried to kill Supervisor!");
    }
}
<commit_msg>Replace test usage of deprecated ind_sample() from rand crate<commit_after>// Copyright (c) 2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Encapsulate running the `hab-sup` executable for tests.

use std::collections::HashSet;
use std::env;
use std::net::TcpListener;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::string::ToString;
use std::sync::Mutex;
use std::thread;
use std::time::Duration;

use hcore::url::BLDR_URL_ENVVAR;
use rand;
use rand::distributions::{Distribution, Uniform};

use super::test_butterfly;

lazy_static! {
    /// Keep track of all TCP ports currently being used by TestSup
    /// instances. Allows us to run tests in parallel without fear of
    /// port conflicts between them.
    static ref CLAIMED_PORTS: Mutex<HashSet<u16>> = {
        Mutex::new(HashSet::new())
    };
}

pub struct TestSup {
    pub hab_root: PathBuf,
    pub origin_name: String,
    pub package_name: String,
    pub service_group: String,
    pub http_port: u16,
    pub butterfly_port: u16,
    pub control_port: u16,
    pub butterfly_client: test_butterfly::Client,
    pub cmd: Command,
    pub process: Option<Child>,
}

/// Return a free TCP port number. We test to see that the system has
/// not already bound the port, while also tracking which ports are
/// being used by other test supervisors that may be running alongside
/// this one.
///
/// Once you receive a port number from this function, you can be
/// reasonably sure that you're the only one that will be using
/// it. There could be a race condition if the machine the tests are
/// running on just happens to claim the same port number for
/// something between the time we check and the time the TestSup
/// claims it. If that happens to you, you should probably buy lottery
/// tickets, though.
///
/// This function will recursively call itself with a decremented
/// value for `tries` if it happens to pick a port that's already in
/// use. Once all tries are used up, it panics! Yay!
fn unclaimed_port(tries: u16) -> u16 {
    if tries == 0 {
        panic!("Couldn't find an unclaimed port for the test Supervisor!")
    }
    let p = random_port();
    match TcpListener::bind(format!("127.0.0.1:{}", p)) {
        Ok(_listener) => {
            // The system hasn't bound it. Now we make sure none of
            // our other tests have bound it.
            let mut ports = CLAIMED_PORTS.lock().unwrap();
            if ports.contains(&p) {
                // Oops, another test is using it, try again
                thread::sleep(Duration::from_millis(500));
                unclaimed_port(tries - 1)
            } else {
                // Nobody was using it. Return the port; the TcpListener
                // that is currently bound to the port will be dropped,
                // thus freeing the port for our use.
                ports.insert(p);
                p
            }
        }
        Err(_) => {
            // port already in use, try again
            unclaimed_port(tries - 1)
        }
    }
}

/// Return a random unprivileged, unregistered TCP port number.
fn random_port() -> u16 {
    // IANA port registrations go to 49151
    let between = Uniform::new_inclusive(49152, ::std::u16::MAX);
    let mut rng = rand::thread_rng();
    between.sample(&mut rng)
}

/// Find an executable relative to the current integration testing
/// executable.
///
/// Thus if the current executable is
///
///    /home/me/habitat/target/debug/deps/compilation-ccaf2f45c24e3840
///
/// and we look for `hab-sup`, we'll find it at
///
///    /home/me/habitat/target/debug/hab-sup
///
fn find_exe<B>(binary_name: B) -> PathBuf
where
    B: AsRef<Path>,
{
    let exe_root = env::current_exe()
        .unwrap()
        .parent()  // deps
        .unwrap()
        .parent()  // debug
        .unwrap()
        .to_path_buf();
    let bin = exe_root.join(binary_name.as_ref());
    assert!(
        bin.exists(),
        format!(
            "Expected to find a {:?} executable at {:?}",
            binary_name.as_ref(),
            bin
        )
    );
    bin
}

/// Return whether or not the tests are being run with the `--nocapture` flag meaning we want to
/// see more output.
fn nocapture_set() -> bool {
    if env::args().any(|arg| arg == "--nocapture") {
        return true;
    } else {
        match env::var("RUST_TEST_NOCAPTURE") {
            Ok(val) => &val != "0",
            Err(_) => false,
        }
    }
}

impl TestSup {
    /// Create a new `TestSup` that will listen on randomly-selected
    /// ports for both gossip and HTTP requests so tests run in
    /// parallel don't step on each other.
    ///
    /// See also `new`.
    pub fn new_with_random_ports<R, O, P, S>(
        fs_root: R,
        origin: O,
        pkg_name: P,
        service_group: S,
    ) -> TestSup
    where
        R: AsRef<Path>,
        O: ToString,
        P: ToString,
        S: ToString,
    {
        // We'll give 10 tries to find a free port number
        let http_port = unclaimed_port(10);
        let butterfly_port = unclaimed_port(10);
        let control_port = unclaimed_port(10);

        TestSup::new(
            fs_root,
            origin,
            pkg_name,
            service_group,
            http_port,
            butterfly_port,
            control_port,
        )
    }

    /// Bundle up a Habitat Supervisor process along with an
    /// associated Butterfly client for injecting new configuration
    /// values. The Supervisor executable is the one that has been
    /// compiled for the current `cargo test` invocation.
    ///
    /// The Supervisor is configured to run a single package for a
    /// test. This package is assumed to have already been installed
    /// relative to `fs_root` (i.e., the `FS_ROOT` environment
    /// variable, which in our tests will be a randomly-named
    /// temporary directory that this Supervisor will view as `/`.).
    ///
    /// A Butterfly client is also created for interacting with this
    /// Supervisor and package. It is properly configured according to
    /// the value provided for `butterfly_port`. To use it, see the
    /// `apply_config` function.
    ///
    /// (No HTTP interaction with the Supervisor is currently called
    /// for, so we don't have a HTTP client.)
    pub fn new<R, O, P, S>(
        fs_root: R,
        origin: O,
        pkg_name: P,
        service_group: S,
        http_port: u16,
        butterfly_port: u16,
        control_port: u16,
    ) -> TestSup
    where
        R: AsRef<Path>,
        O: ToString,
        P: ToString,
        S: ToString,
    {
        let sup_exe = find_exe("hab-sup");
        let launcher_exe = find_exe("hab-launch");

        let mut cmd = Command::new(&launcher_exe);
        let listen_host = "0.0.0.0";

        let origin = origin.to_string();
        let pkg_name = pkg_name.to_string();
        let service_group = service_group.to_string();

        cmd.env(
            "TESTING_FS_ROOT",
            fs_root.as_ref().to_string_lossy().as_ref(),
        ).env("HAB_SUP_BINARY", &sup_exe)
            .env(BLDR_URL_ENVVAR, "http://hab.sup.test")
            .arg("run")
            .arg("--listen-gossip")
            .arg(format!("{}:{}", listen_host, butterfly_port))
            .arg("--listen-http")
            .arg(format!("{}:{}", listen_host, http_port))
            .arg("--listen-ctl")
            .arg(format!("{}:{}", listen_host, control_port))
            // Note: we will have already dropped off the spec files
            // needed to run our test service, so we don't supply a
            // package identifier here
            .stdin(Stdio::null());
        if !nocapture_set() {
            cmd.stdout(Stdio::null());
            cmd.stderr(Stdio::null());
        }

        let bc = test_butterfly::Client::new(&pkg_name, &service_group, butterfly_port);

        TestSup {
            hab_root: fs_root.as_ref().to_path_buf(),
            origin_name: origin,
            package_name: pkg_name,
            service_group: service_group.to_string(),
            http_port: http_port,
            butterfly_port: butterfly_port,
            control_port: control_port,
            butterfly_client: bc,
            cmd: cmd,
            process: None,
        }
    }

    /// Spawn a process actually running the Supervisor.
    pub fn start(&mut self) {
        let child = self.cmd.spawn().expect("Couldn't start the Supervisor!");
        self.process = Some(child);
    }

    /// The equivalent of performing `hab apply` with the given
    /// configuration.
    pub fn apply_config<T>(&mut self, toml_config: T)
    where
        T: ToString,
    {
        self.butterfly_client.apply(toml_config.to_string())
    }
}

// We kill the Supervisor so you don't have to! We also free up the
// ports used by this Supervisor so other tests can use them.
impl Drop for TestSup {
    fn drop(&mut self) {
        let mut ports = CLAIMED_PORTS.lock().unwrap();
        ports.remove(&self.http_port);
        ports.remove(&self.butterfly_port);
        ports.remove(&self.control_port);
        self.process
            .take()
            .expect("No process to kill!")
            .kill()
            .expect("Tried to kill Supervisor!");
    }
}
<|endoftext|>