<commit_before>
fn main() {

    // formatted output
    println!("Hello, Rust!");
    println!("Integer number = {}", 23);
    println!("Float number = {}", 1.012);
    println!("Repeat printing: {0}, {1}, {1}, {3}, {2}", 2, 1, 44, 9);
    println!("Named arguments: {name}, {age}, {occupation}", age="24", name="Ann", occupation="Designer");
    println!("{:08b} XOR {:08b} = {:08b}", 3i8, 32i8, 3i8 ^ 34i8);
    println!("Floating point number {:?}", 2.910210);
    println!("Hex: {:08X}, {:08x}", 131342112, 131342112);
    // TODO find out why this doesn't work:
    //println!("{0:08X}, {0:08x}", 7661);

    // variables
    let integer = 3i32;
    let floatingPoint = 2.1137;

    let mut uInteger = 3u32;
    //integer += 1; error: an immutable variable
    uInteger += 1;

    println!("{}, {}, {}", integer, floatingPoint, uInteger);
    // TODO find out why snake case instead of camel case

    // shadowing
    let first_var = 4;
    {
        let first_var = 51;
        println!("first_var = {}", first_var);
    }
    println!("first_var = {}", first_var);

    // unitialized variables
    let uninit_var;
    //println!("uninit_var = {}", uninit_var); error: use of an init variable
    uninit_var = 231.2;
    println!("already initialized uninit_var = {}", uninit_var);

    // type casting
    let a_integer = 0x40_u8;
    let a_char = a_integer as char;
    println!("char cast from integer {}: {}", a_integer, a_char);

    let a_float = 8.45511;
    println!("integer cast from float {}: {}", a_float, a_float as i32);

    println!("Binary represntation of {}: {:08b}", -1_i8, -1_i8);
    println!("Binary represntation of {}: {:08b}", -2_i8, -2_i8);

    // aliases
    type FloatingPoint32 = f32;
    let a_float_2 : FloatingPoint32 = 3 as f32;
    println!("Type alias: {}, size of var {}", a_float_2, std::mem::size_of_val(&a_float_2));

    // expressions
    let expr = {
        let x = 4.09;
        let y = 1.28;
        println!("x = {}, y = {}", x, y);
        x / y
    };
    println!("expr = {}", expr);

    // if-statement (every if-statement is an expression too, which is really a cool thing)
    let speed = 10;
    let mass = 890;
    let consumption =
        if speed < 0 {
            -1.0
        }
        else if speed == 0 {
            0.8
        }
        else if speed > 0 && speed < 10 {
            18.52
        }
        else if speed >= 10 && speed < 30 {
            let x = mass / speed;
            (x as f32) * 0.05
        }
        else {
            11.71
        };
        println!("Speed = {}, consumption = {}", speed, consumption);

        // loop
        let mut count = 0u32;
        loop {
            println!("loop iteration #{}", count);
            count += 1;
            if count > 9 {
            break;
            }
        }

        count = 0;
        while count < 10 {
            println!("while loop iteration #{}", count);
            count += 1;
        }

        for n in 0u32..10 {
            println!("for loop iteration #{}", n);
        }

        // matching
        let age = 19;
        match age {
            1  => println!("Seems like 1"),
            5  => println!("It's 5"),
            18 => println!("Exactly 18!"),
            _  => println!("I couldn't guess :("),
        }

        match age {
            a if a > 1 && a <= 9  => println!("The age is in range [1,9]"),
            a if a > 9 && a <= 19 => println!("The age is in range [10, 19]"),
            _                     => println!("Nothing, I don't know what range the age is in :("),
        }

        // the previous match is equivalent to the following:
        match age {
            a @ 1...9   => println!("The age is in range [1,9]"),
            a @ 10...19 => println!("The age is in range [10, 19]"),
            a           => println!("Nothing, I don't know what range the age is in :("),
        }

        // if let
        let a_some = Some(0.0008);
        if let Some(i) = a_some {
            println!("a_some is not None, {}", i);
            println!("result {}",  9.0 / i);
        }
        else {
            println!("None");
        }

}
<commit_msg>Progress: while-let, functions, modules<commit_after>
fn main() {

    // formatted output
    println!("Hello, Rust!");
    println!("Integer number = {}", 23);
    println!("Float number = {}", 1.012);
    println!("Repeat printing: {0}, {1}, {1}, {3}, {2}", 2, 1, 44, 9);
    println!("Named arguments: {name}, {age}, {occupation}", age="24", name="Ann", occupation="Designer");
    println!("{:08b} XOR {:08b} = {:08b}", 3i8, 32i8, 3i8 ^ 34i8);
    println!("Floating point number {:?}", 2.910210);
    println!("Hex: {:08X}, {:08x}", 131342112, 131342112);
    // TODO find out why this doesn't work:
    //println!("{0:08X}, {0:08x}", 7661);

    // variables
    let integer = 3i32;
    let floatingPoint = 2.1137;

    let mut uInteger = 3u32;
    //integer += 1; error: an immutable variable
    uInteger += 1;

    println!("{}, {}, {}", integer, floatingPoint, uInteger);
    // TODO find out why snake case instead of camel case

    // shadowing
    let first_var = 4;
    {
        let first_var = 51;
        println!("first_var = {}", first_var);
    }
    println!("first_var = {}", first_var);

    // unitialized variables
    let uninit_var;
    //println!("uninit_var = {}", uninit_var); error: use of an init variable
    uninit_var = 231.2;
    println!("already initialized uninit_var = {}", uninit_var);

    // type casting
    let a_integer = 0x40_u8;
    let a_char = a_integer as char;
    println!("char cast from integer {}: {}", a_integer, a_char);

    let a_float = 8.45511;
    println!("integer cast from float {}: {}", a_float, a_float as i32);

    println!("Binary represntation of {}: {:08b}", -1_i8, -1_i8);
    println!("Binary represntation of {}: {:08b}", -2_i8, -2_i8);

    // aliases
    type FloatingPoint32 = f32;
    let a_float_2 : FloatingPoint32 = 3 as f32;
    println!("Type alias: {}, size of var {}", a_float_2, std::mem::size_of_val(&a_float_2));

    // expressions
    let expr = {
        let x = 4.09;
        let y = 1.28;
        println!("x = {}, y = {}", x, y);
        x / y
    };
    println!("expr = {}", expr);

    // if-statement (every if-statement is an expression too, which is really a cool thing)
    let speed = 10;
    let mass = 890;
    let consumption =
        if speed < 0 {
            -1.0
        }
        else if speed == 0 {
            0.8
        }
        else if speed > 0 && speed < 10 {
            18.52
        }
        else if speed >= 10 && speed < 30 {
            let x = mass / speed;
            (x as f32) * 0.05
        }
        else {
            11.71
        };
        println!("Speed = {}, consumption = {}", speed, consumption);

        // loop
        let mut count = 0u32;
        loop {
            println!("loop iteration #{}", count);
            count += 1;
            if count > 9 {
            break;
            }
        }

        count = 0;
        while count < 10 {
            println!("while loop iteration #{}", count);
            count += 1;
        }

        for n in 0u32..10 {
            println!("for loop iteration #{}", n);
        }

        // matching
        let age = 19;
        match age {
            1  => println!("Seems like 1"),
            5  => println!("It's 5"),
            18 => println!("Exactly 18!"),
            _  => println!("I couldn't guess :("),
        }

        match age {
            a if a > 1 && a <= 9  => println!("The age is in range [1,9]"),
            a if a > 9 && a <= 19 => println!("The age is in range [10, 19]"),
            _                     => println!("Nothing, I don't know what range the age is in :("),
        }

        // the previous match is equivalent to the following:
        match age {
            a @ 1...9   => println!("The age is in range [1,9]"),
            a @ 10...19 => println!("The age is in range [10, 19]"),
            a           => println!("Nothing, I don't know what range the age is in :("),
        }

        // if let
        let a_some = Some(0.0008);
        if let Some(i) = a_some {
            println!("a_some is not None, {}", i);
            println!("result {}",  9.0 / i);
        }
        else {
            println!("None");
        }

        // while let
        let mut a_optional = Some(10);

        while let Some(i) = a_optional {
            if i == 20 {
                println!("while let: i equals 20, leave");
                a_optional = None;
            }
            else {
                a_optional = Some(i + 1);
            }
        }

        // functions, finally!
        let first_integer: i32 = -20;
        let second_integer: i32 = 8;
        let result_of_addition = add_two_integers(first_integer, second_integer);
        println!("{} + {} = {}", first_integer, second_integer, result_of_addition);

        // modules
        mod simple_functions {
            pub fn print_hello_world() {
                println!("Hello, World! (from the simple_functions module)");
            }

            // nested modules
            pub mod simple_math {
                pub fn multiple_two_integers(a: i32, b: i32) -> i32 {
                    a * b
                }
            }
        }

        simple_functions::print_hello_world();
        let mul_result = simple_functions::simple_math::multiple_two_integers(3, 15);
        println!("Modules: multiplication result = {}", mul_result);

        // this doesn't work.. why?
        //use simple_functions::simple_math::multiple_two_integers;
        //println!("Modules: multiplication result 2 = {}", multiple_two_integers(-1, 8));
}

fn add_two_integers(a: i32, b: i32) -> i32 {
    a + b
}
<|endoftext|>


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

use super::{CoreContext, Test};
use crate::Result;
use diem_sdk::{
    client::BlockingClient,
    types::{chain_id::ChainId, LocalAccount},
};

/// The testing interface which defines a test written from the perspective of the Admin of the
/// network. This means that the test will have access to the Root account but do not control any
/// of the validators or full nodes running on the network.
pub trait AdminTest: Test {
    /// Executes the test against the given context.
    fn run<'t>(&self, ctx: &mut AdminContext<'t>) -> Result<()>;
}

#[derive(Debug)]
pub struct AdminContext<'t> {
    core: CoreContext,

    admin_info: AdminInfo<'t>,
}

impl<'t> AdminContext<'t> {
    pub fn new(core: CoreContext, admin_info: AdminInfo<'t>) -> Self {
        Self { core, admin_info }
    }

    pub fn core(&self) -> &CoreContext {
        &self.core
    }

    pub fn client(&self) -> BlockingClient {
        BlockingClient::new(self.admin_info.json_rpc_url)
    }
}

#[derive(Debug)]
pub struct AdminInfo<'t> {
    root_account: &'t mut LocalAccount,
    treasury_compliance_account: &'t mut LocalAccount,

    json_rpc_url: &'t str,
    chain_id: ChainId,
}

impl<'t> AdminInfo<'t> {
    pub fn new(
        root_account: &'t mut LocalAccount,
        treasury_compliance_account: &'t mut LocalAccount,
        json_rpc_url: &'t str,
        chain_id: ChainId,
    ) -> Self {
        Self {
            root_account,
            treasury_compliance_account,
            json_rpc_url,
            chain_id,
        }
    }
}
<commit_msg>forge: have AdminContext take an owned url string<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use super::{CoreContext, Test};
use crate::Result;
use diem_sdk::{
    client::BlockingClient,
    types::{chain_id::ChainId, LocalAccount},
};

/// The testing interface which defines a test written from the perspective of the Admin of the
/// network. This means that the test will have access to the Root account but do not control any
/// of the validators or full nodes running on the network.
pub trait AdminTest: Test {
    /// Executes the test against the given context.
    fn run<'t>(&self, ctx: &mut AdminContext<'t>) -> Result<()>;
}

#[derive(Debug)]
pub struct AdminContext<'t> {
    core: CoreContext,

    admin_info: AdminInfo<'t>,
}

impl<'t> AdminContext<'t> {
    pub fn new(core: CoreContext, admin_info: AdminInfo<'t>) -> Self {
        Self { core, admin_info }
    }

    pub fn core(&self) -> &CoreContext {
        &self.core
    }

    pub fn client(&self) -> BlockingClient {
        BlockingClient::new(&self.admin_info.json_rpc_url)
    }
}

#[derive(Debug)]
pub struct AdminInfo<'t> {
    root_account: &'t mut LocalAccount,
    treasury_compliance_account: &'t mut LocalAccount,

    json_rpc_url: String,
    chain_id: ChainId,
}

impl<'t> AdminInfo<'t> {
    pub fn new(
        root_account: &'t mut LocalAccount,
        treasury_compliance_account: &'t mut LocalAccount,
        json_rpc_url: String,
        chain_id: ChainId,
    ) -> Self {
        Self {
            root_account,
            treasury_compliance_account,
            json_rpc_url,
            chain_id,
        }
    }
}
<|endoftext|>


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

//! Convenience Network API for Diem

pub use crate::protocols::rpc::error::RpcError;
use crate::{
    error::NetworkError,
    peer_manager::{
        ConnectionNotification, ConnectionRequestSender, PeerManagerNotification,
        PeerManagerRequestSender,
    },
    transport::ConnectionMetadata,
    ProtocolId,
};
use bytes::Bytes;
use channel::diem_channel;
use diem_logger::prelude::*;
use diem_types::{network_address::NetworkAddress, PeerId};
use futures::{
    channel::oneshot,
    future,
    stream::{FilterMap, FusedStream, Map, Select, Stream, StreamExt},
    task::{Context, Poll},
};
use pin_project::pin_project;
use serde::{de::DeserializeOwned, Serialize};
use short_hex_str::AsShortHexStr;
use std::{cmp::min, iter::FromIterator, marker::PhantomData, pin::Pin, time::Duration};

use super::wire::handshake::v1::ProtocolIdSet;

pub trait Message: DeserializeOwned + Serialize {}
impl<T: DeserializeOwned + Serialize> Message for T {}

/// Events received by network clients in a validator
///
/// An enumeration of the various types of messages that the network will be sending
/// to its clients. This differs from [`PeerNotification`] since the contents are deserialized
/// into the type `TMessage` over which `Event` is generic. Note that we assume here that for every
/// consumer of this API there's a singleton message type, `TMessage`,  which encapsulates all the
/// messages and RPCs that are received by that consumer.
///
/// [`PeerNotification`]: crate::peer::PeerNotification
#[derive(Debug)]
pub enum Event<TMessage> {
    /// New inbound direct-send message from peer.
    Message(PeerId, TMessage),
    /// New inbound rpc request. The request is fulfilled by sending the
    /// serialized response `Bytes` over the `oneshot::Sender`, where the network
    /// layer will handle sending the response over-the-wire.
    RpcRequest(PeerId, TMessage, oneshot::Sender<Result<Bytes, RpcError>>),
    /// Peer which we have a newly established connection with.
    NewPeer(ConnectionMetadata),
    /// Peer with which we've lost our connection.
    LostPeer(ConnectionMetadata),
}

/// impl PartialEq for simpler testing
impl<TMessage: PartialEq> PartialEq for Event<TMessage> {
    fn eq(&self, other: &Event<TMessage>) -> bool {
        use Event::*;
        match (self, other) {
            (Message(pid1, msg1), Message(pid2, msg2)) => pid1 == pid2 && msg1 == msg2,
            // ignore oneshot::Sender in comparison
            (RpcRequest(pid1, msg1, _), RpcRequest(pid2, msg2, _)) => pid1 == pid2 && msg1 == msg2,
            (NewPeer(metadata1), NewPeer(metadata2)) => metadata1 == metadata2,
            (LostPeer(metadata1), LostPeer(metadata2)) => metadata1 == metadata2,
            _ => false,
        }
    }
}

/// Configuration needed for DiemNet applications to register with the network
/// builder. Supports client-only, service-only, and p2p (both) applications.
// TODO(philiphayes): separate configs for client & server?
#[derive(Clone, Default)]
pub struct AppConfig {
    /// The set of protocols needed for this application.
    pub protocols: ProtocolIdSet,
    /// The config for the inbound message queue from network to the application.
    /// Used for specifying the queue style (e.g. FIFO vs LIFO) and sub-queue max
    /// capacity.
    // TODO(philiphayes): only relevant for services
    // TODO(philiphayes): in the future, use a Service trait here instead?
    pub inbound_queue: Option<diem_channel::Config>,
}

impl AppConfig {
    /// DiemNet client configuration. Requires the set of protocols used by the
    /// client in its requests.
    pub fn client(protocols: impl IntoIterator<Item = ProtocolId>) -> Self {
        Self {
            protocols: ProtocolIdSet::from_iter(protocols),
            inbound_queue: None,
        }
    }

    /// DiemNet service configuration. Requires both the set of protocols this
    /// service can handle and the queue configuration.
    pub fn service(
        protocols: impl IntoIterator<Item = ProtocolId>,
        inbound_queue: diem_channel::Config,
    ) -> Self {
        Self {
            protocols: ProtocolIdSet::from_iter(protocols),
            inbound_queue: Some(inbound_queue),
        }
    }

    /// DiemNet peer-to-peer service configuration. A peer-to-peer service is both
    /// a client and a service.
    pub fn p2p(
        protocols: impl IntoIterator<Item = ProtocolId>,
        inbound_queue: diem_channel::Config,
    ) -> Self {
        Self {
            protocols: ProtocolIdSet::from_iter(protocols),
            inbound_queue: Some(inbound_queue),
        }
    }
}

/// A `Stream` of `Event<TMessage>` from the lower network layer to an upper
/// network application that deserializes inbound network direct-send and rpc
/// messages into `TMessage`. Inbound messages that fail to deserialize are logged
/// and dropped.
///
/// `NetworkEvents` is really just a thin wrapper around a
/// `channel::Receiver<PeerNotification>` that deserializes inbound messages.
#[pin_project]
pub struct NetworkEvents<TMessage> {
    #[pin]
    event_stream: Select<
        FilterMap<
            diem_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
            future::Ready<Option<Event<TMessage>>>,
            fn(PeerManagerNotification) -> future::Ready<Option<Event<TMessage>>>,
        >,
        Map<
            diem_channel::Receiver<PeerId, ConnectionNotification>,
            fn(ConnectionNotification) -> Event<TMessage>,
        >,
    >,
    _marker: PhantomData<TMessage>,
}

/// Trait specifying the signature for `new()` `NetworkEvents`
pub trait NewNetworkEvents {
    fn new(
        peer_mgr_notifs_rx: diem_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
        connection_notifs_rx: diem_channel::Receiver<PeerId, ConnectionNotification>,
    ) -> Self;
}

impl<TMessage: Message> NewNetworkEvents for NetworkEvents<TMessage> {
    fn new(
        peer_mgr_notifs_rx: diem_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
        connection_notifs_rx: diem_channel::Receiver<PeerId, ConnectionNotification>,
    ) -> Self {
        let data_event_stream = peer_mgr_notifs_rx.filter_map(
            peer_mgr_notif_to_event
                as fn(PeerManagerNotification) -> future::Ready<Option<Event<TMessage>>>,
        );
        let control_event_stream = connection_notifs_rx
            .map(control_msg_to_event as fn(ConnectionNotification) -> Event<TMessage>);
        Self {
            event_stream: ::futures::stream::select(data_event_stream, control_event_stream),
            _marker: PhantomData,
        }
    }
}

impl<TMessage> Stream for NetworkEvents<TMessage> {
    type Item = Event<TMessage>;

    fn poll_next(self: Pin<&mut Self>, context: &mut Context) -> Poll<Option<Self::Item>> {
        self.project().event_stream.poll_next(context)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.event_stream.size_hint()
    }
}

/// Deserialize inbound direct send and rpc messages into the application `TMessage`
/// type, logging and dropping messages that fail to deserialize.
fn peer_mgr_notif_to_event<TMessage: Message>(
    notif: PeerManagerNotification,
) -> future::Ready<Option<Event<TMessage>>> {
    let maybe_event = match notif {
        PeerManagerNotification::RecvRpc(peer_id, rpc_req) => {
            match rpc_req.protocol_id.from_bytes(&rpc_req.data) {
                Ok(req_msg) => Some(Event::RpcRequest(peer_id, req_msg, rpc_req.res_tx)),
                Err(err) => {
                    let data = &rpc_req.data;
                    warn!(
                        SecurityEvent::InvalidNetworkEvent,
                        error = ?err,
                        remote_peer_id = peer_id.short_str(),
                        protocol_id = rpc_req.protocol_id,
                        data_prefix = hex::encode(&data[..min(16, data.len())]),
                    );
                    None
                }
            }
        }
        PeerManagerNotification::RecvMessage(peer_id, msg) => {
            match msg.protocol_id.from_bytes(&msg.mdata) {
                Ok(msg) => Some(Event::Message(peer_id, msg)),
                Err(err) => {
                    let data = &msg.mdata;
                    warn!(
                        SecurityEvent::InvalidNetworkEvent,
                        error = ?err,
                        remote_peer_id = peer_id.short_str(),
                        protocol_id = msg.protocol_id,
                        data_prefix = hex::encode(&data[..min(16, data.len())]),
                    );
                    None
                }
            }
        }
    };
    future::ready(maybe_event)
}

fn control_msg_to_event<TMessage>(notif: ConnectionNotification) -> Event<TMessage> {
    match notif {
        ConnectionNotification::NewPeer(metadata, _context) => Event::NewPeer(metadata),
        ConnectionNotification::LostPeer(metadata, _context, _reason) => Event::LostPeer(metadata),
    }
}

impl<TMessage> FusedStream for NetworkEvents<TMessage> {
    fn is_terminated(&self) -> bool {
        self.event_stream.is_terminated()
    }
}

/// `NetworkSender` is the generic interface from upper network applications to
/// the lower network layer. It provides the full API for network applications,
/// including sending direct-send messages, sending rpc requests, as well as
/// dialing or disconnecting from peers and updating the list of accepted public
/// keys.
///
/// `NetworkSender` is in fact a thin wrapper around a `PeerManagerRequestSender`, which in turn is
/// a thin wrapper on `diem_channel::Sender<(PeerId, ProtocolId), PeerManagerRequest>`,
/// mostly focused on providing a more ergonomic API. However, network applications will usually
/// provide their own thin wrapper around `NetworkSender` that narrows the API to the specific
/// interface they need. For instance, `mempool` only requires direct-send functionality so its
/// `MempoolNetworkSender` only exposes a `send_to` function.
///
/// Provide Protobuf wrapper over `[peer_manager::PeerManagerRequestSender]`
#[derive(Clone)]
pub struct NetworkSender<TMessage> {
    peer_mgr_reqs_tx: PeerManagerRequestSender,
    connection_reqs_tx: ConnectionRequestSender,
    _marker: PhantomData<TMessage>,
}

/// Trait specifying the signature for `new()` `NetworkSender`s
pub trait NewNetworkSender {
    fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self;
}

impl<TMessage> NewNetworkSender for NetworkSender<TMessage> {
    fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self {
        Self {
            peer_mgr_reqs_tx,
            connection_reqs_tx,
            _marker: PhantomData,
        }
    }
}

impl<TMessage> NetworkSender<TMessage> {
    /// Request that a given Peer be dialed at the provided `NetworkAddress` and
    /// synchronously wait for the request to be performed.
    pub async fn dial_peer(
        &mut self,
        peer: PeerId,
        addr: NetworkAddress,
    ) -> Result<(), NetworkError> {
        self.connection_reqs_tx.dial_peer(peer, addr).await?;
        Ok(())
    }

    /// Request that a given Peer be disconnected and synchronously wait for the request to be
    /// performed.
    pub async fn disconnect_peer(&mut self, peer: PeerId) -> Result<(), NetworkError> {
        self.connection_reqs_tx.disconnect_peer(peer).await?;
        Ok(())
    }
}

impl<TMessage: Message> NetworkSender<TMessage> {
    /// Send a protobuf message to a single recipient. Provides a wrapper over
    /// `[peer_manager::PeerManagerRequestSender::send_to]`.
    pub fn send_to(
        &mut self,
        recipient: PeerId,
        protocol: ProtocolId,
        message: TMessage,
    ) -> Result<(), NetworkError> {
        let mdata = protocol.to_bytes(&message)?.into();
        self.peer_mgr_reqs_tx.send_to(recipient, protocol, mdata)?;
        Ok(())
    }

    /// Send a protobuf message to a many recipients. Provides a wrapper over
    /// `[peer_manager::PeerManagerRequestSender::send_to_many]`.
    pub fn send_to_many(
        &mut self,
        recipients: impl Iterator<Item = PeerId>,
        protocol: ProtocolId,
        message: TMessage,
    ) -> Result<(), NetworkError> {
        // Serialize message.
        let mdata = protocol.to_bytes(&message)?.into();
        self.peer_mgr_reqs_tx
            .send_to_many(recipients, protocol, mdata)?;
        Ok(())
    }

    /// Send a protobuf rpc request to a single recipient while handling
    /// serialization and deserialization of the request and response respectively.
    /// Assumes that the request and response both have the same message type.
    pub async fn send_rpc(
        &mut self,
        recipient: PeerId,
        protocol: ProtocolId,
        req_msg: TMessage,
        timeout: Duration,
    ) -> Result<TMessage, RpcError> {
        // serialize request
        let req_data = protocol.to_bytes(&req_msg)?.into();
        let res_data = self
            .peer_mgr_reqs_tx
            .send_rpc(recipient, protocol, req_data, timeout)
            .await?;
        let res_msg: TMessage = protocol.from_bytes(&res_data)?;
        Ok(res_msg)
    }
}
<commit_msg>[network] Make trait for ApplicationNetworkSenders<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Convenience Network API for Diem

pub use crate::protocols::rpc::error::RpcError;
use crate::{
    error::NetworkError,
    peer_manager::{
        ConnectionNotification, ConnectionRequestSender, PeerManagerNotification,
        PeerManagerRequestSender,
    },
    transport::ConnectionMetadata,
    ProtocolId,
};
use async_trait::async_trait;
use bytes::Bytes;
use channel::diem_channel;
use diem_logger::prelude::*;
use diem_types::{network_address::NetworkAddress, PeerId};
use futures::{
    channel::oneshot,
    future,
    stream::{FilterMap, FusedStream, Map, Select, Stream, StreamExt},
    task::{Context, Poll},
};
use pin_project::pin_project;
use serde::{de::DeserializeOwned, Serialize};
use short_hex_str::AsShortHexStr;
use std::{cmp::min, iter::FromIterator, marker::PhantomData, pin::Pin, time::Duration};

use super::wire::handshake::v1::ProtocolIdSet;

pub trait Message: DeserializeOwned + Serialize {}
impl<T: DeserializeOwned + Serialize> Message for T {}

/// Events received by network clients in a validator
///
/// An enumeration of the various types of messages that the network will be sending
/// to its clients. This differs from [`PeerNotification`] since the contents are deserialized
/// into the type `TMessage` over which `Event` is generic. Note that we assume here that for every
/// consumer of this API there's a singleton message type, `TMessage`,  which encapsulates all the
/// messages and RPCs that are received by that consumer.
///
/// [`PeerNotification`]: crate::peer::PeerNotification
#[derive(Debug)]
pub enum Event<TMessage> {
    /// New inbound direct-send message from peer.
    Message(PeerId, TMessage),
    /// New inbound rpc request. The request is fulfilled by sending the
    /// serialized response `Bytes` over the `oneshot::Sender`, where the network
    /// layer will handle sending the response over-the-wire.
    RpcRequest(PeerId, TMessage, oneshot::Sender<Result<Bytes, RpcError>>),
    /// Peer which we have a newly established connection with.
    NewPeer(ConnectionMetadata),
    /// Peer with which we've lost our connection.
    LostPeer(ConnectionMetadata),
}

/// impl PartialEq for simpler testing
impl<TMessage: PartialEq> PartialEq for Event<TMessage> {
    fn eq(&self, other: &Event<TMessage>) -> bool {
        use Event::*;
        match (self, other) {
            (Message(pid1, msg1), Message(pid2, msg2)) => pid1 == pid2 && msg1 == msg2,
            // ignore oneshot::Sender in comparison
            (RpcRequest(pid1, msg1, _), RpcRequest(pid2, msg2, _)) => pid1 == pid2 && msg1 == msg2,
            (NewPeer(metadata1), NewPeer(metadata2)) => metadata1 == metadata2,
            (LostPeer(metadata1), LostPeer(metadata2)) => metadata1 == metadata2,
            _ => false,
        }
    }
}

/// Configuration needed for DiemNet applications to register with the network
/// builder. Supports client-only, service-only, and p2p (both) applications.
// TODO(philiphayes): separate configs for client & server?
#[derive(Clone, Default)]
pub struct AppConfig {
    /// The set of protocols needed for this application.
    pub protocols: ProtocolIdSet,
    /// The config for the inbound message queue from network to the application.
    /// Used for specifying the queue style (e.g. FIFO vs LIFO) and sub-queue max
    /// capacity.
    // TODO(philiphayes): only relevant for services
    // TODO(philiphayes): in the future, use a Service trait here instead?
    pub inbound_queue: Option<diem_channel::Config>,
}

impl AppConfig {
    /// DiemNet client configuration. Requires the set of protocols used by the
    /// client in its requests.
    pub fn client(protocols: impl IntoIterator<Item = ProtocolId>) -> Self {
        Self {
            protocols: ProtocolIdSet::from_iter(protocols),
            inbound_queue: None,
        }
    }

    /// DiemNet service configuration. Requires both the set of protocols this
    /// service can handle and the queue configuration.
    pub fn service(
        protocols: impl IntoIterator<Item = ProtocolId>,
        inbound_queue: diem_channel::Config,
    ) -> Self {
        Self {
            protocols: ProtocolIdSet::from_iter(protocols),
            inbound_queue: Some(inbound_queue),
        }
    }

    /// DiemNet peer-to-peer service configuration. A peer-to-peer service is both
    /// a client and a service.
    pub fn p2p(
        protocols: impl IntoIterator<Item = ProtocolId>,
        inbound_queue: diem_channel::Config,
    ) -> Self {
        Self {
            protocols: ProtocolIdSet::from_iter(protocols),
            inbound_queue: Some(inbound_queue),
        }
    }
}

/// A `Stream` of `Event<TMessage>` from the lower network layer to an upper
/// network application that deserializes inbound network direct-send and rpc
/// messages into `TMessage`. Inbound messages that fail to deserialize are logged
/// and dropped.
///
/// `NetworkEvents` is really just a thin wrapper around a
/// `channel::Receiver<PeerNotification>` that deserializes inbound messages.
#[pin_project]
pub struct NetworkEvents<TMessage> {
    #[pin]
    event_stream: Select<
        FilterMap<
            diem_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
            future::Ready<Option<Event<TMessage>>>,
            fn(PeerManagerNotification) -> future::Ready<Option<Event<TMessage>>>,
        >,
        Map<
            diem_channel::Receiver<PeerId, ConnectionNotification>,
            fn(ConnectionNotification) -> Event<TMessage>,
        >,
    >,
    _marker: PhantomData<TMessage>,
}

/// Trait specifying the signature for `new()` `NetworkEvents`
pub trait NewNetworkEvents {
    fn new(
        peer_mgr_notifs_rx: diem_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
        connection_notifs_rx: diem_channel::Receiver<PeerId, ConnectionNotification>,
    ) -> Self;
}

impl<TMessage: Message> NewNetworkEvents for NetworkEvents<TMessage> {
    fn new(
        peer_mgr_notifs_rx: diem_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
        connection_notifs_rx: diem_channel::Receiver<PeerId, ConnectionNotification>,
    ) -> Self {
        let data_event_stream = peer_mgr_notifs_rx.filter_map(
            peer_mgr_notif_to_event
                as fn(PeerManagerNotification) -> future::Ready<Option<Event<TMessage>>>,
        );
        let control_event_stream = connection_notifs_rx
            .map(control_msg_to_event as fn(ConnectionNotification) -> Event<TMessage>);
        Self {
            event_stream: ::futures::stream::select(data_event_stream, control_event_stream),
            _marker: PhantomData,
        }
    }
}

impl<TMessage> Stream for NetworkEvents<TMessage> {
    type Item = Event<TMessage>;

    fn poll_next(self: Pin<&mut Self>, context: &mut Context) -> Poll<Option<Self::Item>> {
        self.project().event_stream.poll_next(context)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.event_stream.size_hint()
    }
}

/// Deserialize inbound direct send and rpc messages into the application `TMessage`
/// type, logging and dropping messages that fail to deserialize.
fn peer_mgr_notif_to_event<TMessage: Message>(
    notif: PeerManagerNotification,
) -> future::Ready<Option<Event<TMessage>>> {
    let maybe_event = match notif {
        PeerManagerNotification::RecvRpc(peer_id, rpc_req) => {
            match rpc_req.protocol_id.from_bytes(&rpc_req.data) {
                Ok(req_msg) => Some(Event::RpcRequest(peer_id, req_msg, rpc_req.res_tx)),
                Err(err) => {
                    let data = &rpc_req.data;
                    warn!(
                        SecurityEvent::InvalidNetworkEvent,
                        error = ?err,
                        remote_peer_id = peer_id.short_str(),
                        protocol_id = rpc_req.protocol_id,
                        data_prefix = hex::encode(&data[..min(16, data.len())]),
                    );
                    None
                }
            }
        }
        PeerManagerNotification::RecvMessage(peer_id, msg) => {
            match msg.protocol_id.from_bytes(&msg.mdata) {
                Ok(msg) => Some(Event::Message(peer_id, msg)),
                Err(err) => {
                    let data = &msg.mdata;
                    warn!(
                        SecurityEvent::InvalidNetworkEvent,
                        error = ?err,
                        remote_peer_id = peer_id.short_str(),
                        protocol_id = msg.protocol_id,
                        data_prefix = hex::encode(&data[..min(16, data.len())]),
                    );
                    None
                }
            }
        }
    };
    future::ready(maybe_event)
}

fn control_msg_to_event<TMessage>(notif: ConnectionNotification) -> Event<TMessage> {
    match notif {
        ConnectionNotification::NewPeer(metadata, _context) => Event::NewPeer(metadata),
        ConnectionNotification::LostPeer(metadata, _context, _reason) => Event::LostPeer(metadata),
    }
}

impl<TMessage> FusedStream for NetworkEvents<TMessage> {
    fn is_terminated(&self) -> bool {
        self.event_stream.is_terminated()
    }
}

/// `NetworkSender` is the generic interface from upper network applications to
/// the lower network layer. It provides the full API for network applications,
/// including sending direct-send messages, sending rpc requests, as well as
/// dialing or disconnecting from peers and updating the list of accepted public
/// keys.
///
/// `NetworkSender` is in fact a thin wrapper around a `PeerManagerRequestSender`, which in turn is
/// a thin wrapper on `diem_channel::Sender<(PeerId, ProtocolId), PeerManagerRequest>`,
/// mostly focused on providing a more ergonomic API. However, network applications will usually
/// provide their own thin wrapper around `NetworkSender` that narrows the API to the specific
/// interface they need. For instance, `mempool` only requires direct-send functionality so its
/// `MempoolNetworkSender` only exposes a `send_to` function.
///
/// Provide Protobuf wrapper over `[peer_manager::PeerManagerRequestSender]`
#[derive(Clone)]
pub struct NetworkSender<TMessage> {
    peer_mgr_reqs_tx: PeerManagerRequestSender,
    connection_reqs_tx: ConnectionRequestSender,
    _marker: PhantomData<TMessage>,
}

/// Trait specifying the signature for `new()` `NetworkSender`s
pub trait NewNetworkSender {
    fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self;
}

impl<TMessage> NewNetworkSender for NetworkSender<TMessage> {
    fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self {
        Self {
            peer_mgr_reqs_tx,
            connection_reqs_tx,
            _marker: PhantomData,
        }
    }
}

impl<TMessage> NetworkSender<TMessage> {
    /// Request that a given Peer be dialed at the provided `NetworkAddress` and
    /// synchronously wait for the request to be performed.
    pub async fn dial_peer(
        &mut self,
        peer: PeerId,
        addr: NetworkAddress,
    ) -> Result<(), NetworkError> {
        self.connection_reqs_tx.dial_peer(peer, addr).await?;
        Ok(())
    }

    /// Request that a given Peer be disconnected and synchronously wait for the request to be
    /// performed.
    pub async fn disconnect_peer(&mut self, peer: PeerId) -> Result<(), NetworkError> {
        self.connection_reqs_tx.disconnect_peer(peer).await?;
        Ok(())
    }
}

impl<TMessage: Message> NetworkSender<TMessage> {
    /// Send a protobuf message to a single recipient. Provides a wrapper over
    /// `[peer_manager::PeerManagerRequestSender::send_to]`.
    pub fn send_to(
        &mut self,
        recipient: PeerId,
        protocol: ProtocolId,
        message: TMessage,
    ) -> Result<(), NetworkError> {
        let mdata = protocol.to_bytes(&message)?.into();
        self.peer_mgr_reqs_tx.send_to(recipient, protocol, mdata)?;
        Ok(())
    }

    /// Send a protobuf message to a many recipients. Provides a wrapper over
    /// `[peer_manager::PeerManagerRequestSender::send_to_many]`.
    pub fn send_to_many(
        &mut self,
        recipients: impl Iterator<Item = PeerId>,
        protocol: ProtocolId,
        message: TMessage,
    ) -> Result<(), NetworkError> {
        // Serialize message.
        let mdata = protocol.to_bytes(&message)?.into();
        self.peer_mgr_reqs_tx
            .send_to_many(recipients, protocol, mdata)?;
        Ok(())
    }

    /// Send a protobuf rpc request to a single recipient while handling
    /// serialization and deserialization of the request and response respectively.
    /// Assumes that the request and response both have the same message type.
    pub async fn send_rpc(
        &mut self,
        recipient: PeerId,
        protocol: ProtocolId,
        req_msg: TMessage,
        timeout: Duration,
    ) -> Result<TMessage, RpcError> {
        // serialize request
        let req_data = protocol.to_bytes(&req_msg)?.into();
        let res_data = self
            .peer_mgr_reqs_tx
            .send_rpc(recipient, protocol, req_data, timeout)
            .await?;
        let res_msg: TMessage = protocol.from_bytes(&res_data)?;
        Ok(res_msg)
    }
}

/// A simplified version of `NetworkSender` that doesn't use `ProtocolId` in the input
/// It was already being implemented for every application, but is now standardized
#[async_trait]
pub trait ApplicationNetworkSender<TMessage: Send>: Clone {
    fn send_to(&mut self, recipient: PeerId, message: TMessage) -> Result<(), NetworkError>;

    fn send_to_many(
        &mut self,
        recipients: impl Iterator<Item = PeerId>,
        message: TMessage,
    ) -> Result<(), NetworkError>;

    async fn send_rpc(
        &mut self,
        recipient: PeerId,
        req_msg: TMessage,
        timeout: Duration,
    ) -> Result<TMessage, RpcError>;
}
<|endoftext|>


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

//! Network API for [`Consensus`](/consensus/index.html) and [`Mempool`](/mempool/index.html)

pub use crate::protocols::rpc::error::RpcError;
use crate::{
    error::NetworkError,
    peer_manager::{self, PeerManagerNotification, PeerManagerRequestSender},
    utils::MessageExt,
    ProtocolId,
};
use bytes::Bytes;
use channel::libra_channel;
use futures::{
    channel::oneshot,
    stream::{FusedStream, Map, Select, Stream, StreamExt},
    task::{Context, Poll},
};
use parity_multiaddr::Multiaddr;
use pin_project::pin_project;
use prost::Message;
use std::{default::Default, marker::PhantomData, pin::Pin, time::Duration};

pub mod consensus;
pub mod discovery;
pub mod health_checker;
pub mod mempool;
pub mod network_builder;
pub mod state_synchronizer;

#[cfg(test)]
mod test;

// Public re-exports
pub use consensus::{
    ConsensusNetworkEvents, ConsensusNetworkSender, CONSENSUS_DIRECT_SEND_PROTOCOL,
    CONSENSUS_RPC_PROTOCOL,
};
pub use discovery::{
    DiscoveryNetworkEvents, DiscoveryNetworkSender, DISCOVERY_DIRECT_SEND_PROTOCOL,
};
pub use health_checker::{
    HealthCheckerNetworkEvents, HealthCheckerNetworkSender, HEALTH_CHECKER_RPC_PROTOCOL,
};
use libra_types::PeerId;
pub use mempool::{MempoolNetworkEvents, MempoolNetworkSender, MEMPOOL_DIRECT_SEND_PROTOCOL};
pub use state_synchronizer::{
    StateSynchronizerEvents, StateSynchronizerSender, STATE_SYNCHRONIZER_DIRECT_SEND_PROTOCOL,
};

/// Events received by network clients in a validator
///
/// An enumeration of the various types of messages that the network will be sending
/// to its clients. This differs from [`NetworkNotification`] since the contents are deserialized
/// into the type `TMessage` over which `Event` is generic. Note that we assume here that for every
/// consumer of this API there's a singleton message type, `TMessage`,  which encapsulates all the
/// messages and RPCs that are received by that consumer.
///
/// [`NetworkNotification`]: crate::interface::NetworkNotification
#[derive(Debug)]
pub enum Event<TMessage> {
    /// New inbound direct-send message from peer.
    Message((PeerId, TMessage)),
    /// New inbound rpc request. The request is fulfilled by sending the
    /// serialized response `Bytes` over the `oneshot::Sender`, where the network
    /// layer will handle sending the response over-the-wire.
    RpcRequest((PeerId, TMessage, oneshot::Sender<Result<Bytes, RpcError>>)),
    /// Peer which we have a newly established connection with.
    NewPeer(PeerId),
    /// Peer with which we've lost our connection.
    LostPeer(PeerId),
}

/// impl PartialEq for simpler testing
impl<TMessage: PartialEq> PartialEq for Event<TMessage> {
    fn eq(&self, other: &Event<TMessage>) -> bool {
        use Event::*;
        match (self, other) {
            (Message((pid1, msg1)), Message((pid2, msg2))) => pid1 == pid2 && msg1 == msg2,
            // ignore oneshot::Sender in comparison
            (RpcRequest((pid1, msg1, _)), RpcRequest((pid2, msg2, _))) => {
                pid1 == pid2 && msg1 == msg2
            }
            (NewPeer(pid1), NewPeer(pid2)) => pid1 == pid2,
            (LostPeer(pid1), LostPeer(pid2)) => pid1 == pid2,
            _ => false,
        }
    }
}

/// A `Stream` of `Event<TMessage>` from the lower network layer to an upper
/// network application that deserializes inbound network direct-send and rpc
/// messages into `TMessage`, which is some protobuf format implementing
/// the `prost::Message` trait.
///
/// `NetworkEvents` is really just a thin wrapper around a
/// `channel::Receiver<NetworkNotification>` that deserializes inbound messages.
#[pin_project]
pub struct NetworkEvents<TMessage: Message + Default> {
    #[pin]
    event_stream: Select<
        Map<
            libra_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
            fn(PeerManagerNotification) -> Result<Event<TMessage>, NetworkError>,
        >,
        Map<
            libra_channel::Receiver<PeerId, peer_manager::ConnectionStatusNotification>,
            fn(peer_manager::ConnectionStatusNotification) -> Result<Event<TMessage>, NetworkError>,
        >,
    >,
    _marker: PhantomData<TMessage>,
}

impl<TMessage: Message + Default> NetworkEvents<TMessage> {
    pub fn new(
        peer_mgr_notifs_rx: libra_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
        control_notifs_rx: libra_channel::Receiver<
            PeerId,
            peer_manager::ConnectionStatusNotification,
        >,
    ) -> Self {
        let data_event_stream = peer_mgr_notifs_rx.map(
            peer_mgr_notif_to_event
                as fn(PeerManagerNotification) -> Result<Event<TMessage>, NetworkError>,
        );
        let control_event_stream = control_notifs_rx.map(
            control_msg_to_event
                as fn(
                    peer_manager::ConnectionStatusNotification,
                ) -> Result<Event<TMessage>, NetworkError>,
        );
        Self {
            event_stream: ::futures::stream::select(data_event_stream, control_event_stream),
            _marker: PhantomData,
        }
    }
}

impl<TMessage: Message + Default> Stream for NetworkEvents<TMessage> {
    type Item = Result<Event<TMessage>, NetworkError>;

    fn poll_next(self: Pin<&mut Self>, context: &mut Context) -> Poll<Option<Self::Item>> {
        self.project().event_stream.poll_next(context)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.event_stream.size_hint()
    }
}

fn peer_mgr_notif_to_event<TMessage>(
    notif: PeerManagerNotification,
) -> Result<Event<TMessage>, NetworkError>
where
    TMessage: Message + Default,
{
    match notif {
        PeerManagerNotification::RecvRpc(peer_id, rpc_req) => {
            let req_msg = TMessage::decode(rpc_req.data.as_ref())?;
            Ok(Event::RpcRequest((peer_id, req_msg, rpc_req.res_tx)))
        }
        PeerManagerNotification::RecvMessage(peer_id, msg) => {
            let msg = TMessage::decode(msg.mdata.as_ref())?;
            Ok(Event::Message((peer_id, msg)))
        }
    }
}

fn control_msg_to_event<TMessage>(
    notif: peer_manager::ConnectionStatusNotification,
) -> Result<Event<TMessage>, NetworkError> {
    match notif {
        peer_manager::ConnectionStatusNotification::NewPeer(peer_id, _addr) => {
            Ok(Event::NewPeer(peer_id))
        }
        peer_manager::ConnectionStatusNotification::LostPeer(peer_id, _addr, _reason) => {
            Ok(Event::LostPeer(peer_id))
        }
    }
}

impl<TMessage: Message + Default> FusedStream for NetworkEvents<TMessage> {
    fn is_terminated(&self) -> bool {
        self.event_stream.is_terminated()
    }
}

/// `NetworkSender` is the generic interface from upper network applications to
/// the lower network layer. It provides the full API for network applications,
/// including sending direct-send messages, sending rpc requests, as well as
/// dialing or disconnecting from peers and updating the list of accepted public
/// keys.
///
/// `NetworkSender` is in fact a thin wrapper around a `PeerManagerRequestSender`, which in turn is
/// a thin wrapper on `libra_channel::Sender<(PeerId, ProtocolId), PeerManagerRequest>`,
/// mostly focused on providing a more ergonomic API. However, network applications will usually
/// provide their own thin wrapper around `NetworkSender` that narrows the API to the specific
/// interface they need. For instance, `mempool` only requires direct-send functionality so its
/// `MempoolNetworkSender` only exposes a `send_to` function.
///
/// Provide Protobuf wrapper over `[peer_manager::PeerManagerRequestSender]`
/// The `TMessage` generic is a protobuf message type (`prost::Message`).
#[derive(Clone)]
pub struct NetworkSender<TMessage: Message + Default> {
    inner: PeerManagerRequestSender,
    _marker: PhantomData<TMessage>,
}

impl<TMessage: Message + Default> NetworkSender<TMessage> {
    pub fn new(inner: PeerManagerRequestSender) -> Self {
        Self {
            inner,
            _marker: PhantomData,
        }
    }

    /// Unwrap the `NetworkSender` into the underlying `PeerManagerRequestSender`.
    pub fn into_inner(self) -> PeerManagerRequestSender {
        self.inner
    }

    /// Get a mutable reference to the underlying `PeerManagerRequestSender`.
    pub fn get_mut(&mut self) -> &mut PeerManagerRequestSender {
        &mut self.inner
    }

    pub async fn dial_peer(&mut self, peer: PeerId, addr: Multiaddr) -> Result<(), NetworkError> {
        self.inner.dial_peer(peer, addr).await?;
        Ok(())
    }

    pub async fn disconnect_peer(&mut self, peer: PeerId) -> Result<(), NetworkError> {
        self.inner.disconnect_peer(peer).await?;
        Ok(())
    }

    /// Send a protobuf message to a single recipient. Provides a wrapper over
    /// `[peer_manager::PeerManagerRequestSender::send_to]`.
    pub fn send_to(
        &mut self,
        recipient: PeerId,
        protocol: ProtocolId,
        message: TMessage,
    ) -> Result<(), NetworkError> {
        // Serialize message.
        let mdata = message.to_bytes().unwrap();
        self.inner.send_to(recipient, protocol, mdata)?;
        Ok(())
    }

    /// Send a protobuf message to a many recipients. Provides a wrapper over
    /// `[peer_manager::PeerManagerRequestSender::send_to_many]`.
    pub fn send_to_many(
        &mut self,
        recipients: impl Iterator<Item = PeerId>,
        protocol: ProtocolId,
        message: TMessage,
    ) -> Result<(), NetworkError> {
        // Serialize message.
        let mdata = message.to_bytes().unwrap();
        self.inner.send_to_many(recipients, protocol, mdata)?;
        Ok(())
    }

    /// Send a protobuf rpc request to a single recipient while handling
    /// serialization and deserialization of the request and response respectively.
    /// Assumes that the request and response both have the same message type.
    pub async fn unary_rpc(
        &mut self,
        recipient: PeerId,
        protocol: ProtocolId,
        req_msg: TMessage,
        timeout: Duration,
    ) -> Result<TMessage, RpcError> {
        // serialize request
        let req_data = req_msg.to_bytes().unwrap();
        let res_data = self
            .inner
            .unary_rpc(recipient, protocol, req_data, timeout)
            .await?;
        let res_msg = TMessage::decode(res_data.as_ref())?;
        Ok(res_msg)
    }
}
<commit_msg>[network] Remove unnecessary trait bounds in NetworkSender/Events impls<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Network API for [`Consensus`](/consensus/index.html) and [`Mempool`](/mempool/index.html)

pub use crate::protocols::rpc::error::RpcError;
use crate::{
    error::NetworkError,
    peer_manager::{self, PeerManagerNotification, PeerManagerRequestSender},
    utils::MessageExt,
    ProtocolId,
};
use bytes::Bytes;
use channel::libra_channel;
use futures::{
    channel::oneshot,
    stream::{FusedStream, Map, Select, Stream, StreamExt},
    task::{Context, Poll},
};
use parity_multiaddr::Multiaddr;
use pin_project::pin_project;
use prost::Message;
use std::{default::Default, marker::PhantomData, pin::Pin, time::Duration};

pub mod consensus;
pub mod discovery;
pub mod health_checker;
pub mod mempool;
pub mod network_builder;
pub mod state_synchronizer;

#[cfg(test)]
mod test;

// Public re-exports
pub use consensus::{
    ConsensusNetworkEvents, ConsensusNetworkSender, CONSENSUS_DIRECT_SEND_PROTOCOL,
    CONSENSUS_RPC_PROTOCOL,
};
pub use discovery::{
    DiscoveryNetworkEvents, DiscoveryNetworkSender, DISCOVERY_DIRECT_SEND_PROTOCOL,
};
pub use health_checker::{
    HealthCheckerNetworkEvents, HealthCheckerNetworkSender, HEALTH_CHECKER_RPC_PROTOCOL,
};
use libra_types::PeerId;
pub use mempool::{MempoolNetworkEvents, MempoolNetworkSender, MEMPOOL_DIRECT_SEND_PROTOCOL};
pub use state_synchronizer::{
    StateSynchronizerEvents, StateSynchronizerSender, STATE_SYNCHRONIZER_DIRECT_SEND_PROTOCOL,
};

/// Events received by network clients in a validator
///
/// An enumeration of the various types of messages that the network will be sending
/// to its clients. This differs from [`NetworkNotification`] since the contents are deserialized
/// into the type `TMessage` over which `Event` is generic. Note that we assume here that for every
/// consumer of this API there's a singleton message type, `TMessage`,  which encapsulates all the
/// messages and RPCs that are received by that consumer.
///
/// [`NetworkNotification`]: crate::interface::NetworkNotification
#[derive(Debug)]
pub enum Event<TMessage> {
    /// New inbound direct-send message from peer.
    Message((PeerId, TMessage)),
    /// New inbound rpc request. The request is fulfilled by sending the
    /// serialized response `Bytes` over the `oneshot::Sender`, where the network
    /// layer will handle sending the response over-the-wire.
    RpcRequest((PeerId, TMessage, oneshot::Sender<Result<Bytes, RpcError>>)),
    /// Peer which we have a newly established connection with.
    NewPeer(PeerId),
    /// Peer with which we've lost our connection.
    LostPeer(PeerId),
}

/// impl PartialEq for simpler testing
impl<TMessage: PartialEq> PartialEq for Event<TMessage> {
    fn eq(&self, other: &Event<TMessage>) -> bool {
        use Event::*;
        match (self, other) {
            (Message((pid1, msg1)), Message((pid2, msg2))) => pid1 == pid2 && msg1 == msg2,
            // ignore oneshot::Sender in comparison
            (RpcRequest((pid1, msg1, _)), RpcRequest((pid2, msg2, _))) => {
                pid1 == pid2 && msg1 == msg2
            }
            (NewPeer(pid1), NewPeer(pid2)) => pid1 == pid2,
            (LostPeer(pid1), LostPeer(pid2)) => pid1 == pid2,
            _ => false,
        }
    }
}

/// A `Stream` of `Event<TMessage>` from the lower network layer to an upper
/// network application that deserializes inbound network direct-send and rpc
/// messages into `TMessage`, which is some protobuf format implementing
/// the `prost::Message` trait.
///
/// `NetworkEvents` is really just a thin wrapper around a
/// `channel::Receiver<NetworkNotification>` that deserializes inbound messages.
#[pin_project]
pub struct NetworkEvents<TMessage> {
    #[pin]
    event_stream: Select<
        Map<
            libra_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
            fn(PeerManagerNotification) -> Result<Event<TMessage>, NetworkError>,
        >,
        Map<
            libra_channel::Receiver<PeerId, peer_manager::ConnectionStatusNotification>,
            fn(peer_manager::ConnectionStatusNotification) -> Result<Event<TMessage>, NetworkError>,
        >,
    >,
    _marker: PhantomData<TMessage>,
}

impl<TMessage: Message + Default> NetworkEvents<TMessage> {
    pub fn new(
        peer_mgr_notifs_rx: libra_channel::Receiver<(PeerId, ProtocolId), PeerManagerNotification>,
        control_notifs_rx: libra_channel::Receiver<
            PeerId,
            peer_manager::ConnectionStatusNotification,
        >,
    ) -> Self {
        let data_event_stream = peer_mgr_notifs_rx.map(
            peer_mgr_notif_to_event
                as fn(PeerManagerNotification) -> Result<Event<TMessage>, NetworkError>,
        );
        let control_event_stream = control_notifs_rx.map(
            control_msg_to_event
                as fn(
                    peer_manager::ConnectionStatusNotification,
                ) -> Result<Event<TMessage>, NetworkError>,
        );
        Self {
            event_stream: ::futures::stream::select(data_event_stream, control_event_stream),
            _marker: PhantomData,
        }
    }
}

impl<TMessage> Stream for NetworkEvents<TMessage> {
    type Item = Result<Event<TMessage>, NetworkError>;

    fn poll_next(self: Pin<&mut Self>, context: &mut Context) -> Poll<Option<Self::Item>> {
        self.project().event_stream.poll_next(context)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.event_stream.size_hint()
    }
}

fn peer_mgr_notif_to_event<TMessage>(
    notif: PeerManagerNotification,
) -> Result<Event<TMessage>, NetworkError>
where
    TMessage: Message + Default,
{
    match notif {
        PeerManagerNotification::RecvRpc(peer_id, rpc_req) => {
            let req_msg = TMessage::decode(rpc_req.data.as_ref())?;
            Ok(Event::RpcRequest((peer_id, req_msg, rpc_req.res_tx)))
        }
        PeerManagerNotification::RecvMessage(peer_id, msg) => {
            let msg = TMessage::decode(msg.mdata.as_ref())?;
            Ok(Event::Message((peer_id, msg)))
        }
    }
}

fn control_msg_to_event<TMessage>(
    notif: peer_manager::ConnectionStatusNotification,
) -> Result<Event<TMessage>, NetworkError> {
    match notif {
        peer_manager::ConnectionStatusNotification::NewPeer(peer_id, _addr) => {
            Ok(Event::NewPeer(peer_id))
        }
        peer_manager::ConnectionStatusNotification::LostPeer(peer_id, _addr, _reason) => {
            Ok(Event::LostPeer(peer_id))
        }
    }
}

impl<TMessage> FusedStream for NetworkEvents<TMessage> {
    fn is_terminated(&self) -> bool {
        self.event_stream.is_terminated()
    }
}

/// `NetworkSender` is the generic interface from upper network applications to
/// the lower network layer. It provides the full API for network applications,
/// including sending direct-send messages, sending rpc requests, as well as
/// dialing or disconnecting from peers and updating the list of accepted public
/// keys.
///
/// `NetworkSender` is in fact a thin wrapper around a `PeerManagerRequestSender`, which in turn is
/// a thin wrapper on `libra_channel::Sender<(PeerId, ProtocolId), PeerManagerRequest>`,
/// mostly focused on providing a more ergonomic API. However, network applications will usually
/// provide their own thin wrapper around `NetworkSender` that narrows the API to the specific
/// interface they need. For instance, `mempool` only requires direct-send functionality so its
/// `MempoolNetworkSender` only exposes a `send_to` function.
///
/// Provide Protobuf wrapper over `[peer_manager::PeerManagerRequestSender]`
/// The `TMessage` generic is a protobuf message type (`prost::Message`).
#[derive(Clone)]
pub struct NetworkSender<TMessage> {
    inner: PeerManagerRequestSender,
    _marker: PhantomData<TMessage>,
}

impl<TMessage> NetworkSender<TMessage> {
    pub fn new(inner: PeerManagerRequestSender) -> Self {
        Self {
            inner,
            _marker: PhantomData,
        }
    }

    /// Unwrap the `NetworkSender` into the underlying `PeerManagerRequestSender`.
    pub fn into_inner(self) -> PeerManagerRequestSender {
        self.inner
    }

    /// Get a mutable reference to the underlying `PeerManagerRequestSender`.
    pub fn get_mut(&mut self) -> &mut PeerManagerRequestSender {
        &mut self.inner
    }

    pub async fn dial_peer(&mut self, peer: PeerId, addr: Multiaddr) -> Result<(), NetworkError> {
        self.inner.dial_peer(peer, addr).await?;
        Ok(())
    }

    pub async fn disconnect_peer(&mut self, peer: PeerId) -> Result<(), NetworkError> {
        self.inner.disconnect_peer(peer).await?;
        Ok(())
    }
}

impl<TMessage: Message + Default> NetworkSender<TMessage> {
    /// Send a protobuf message to a single recipient. Provides a wrapper over
    /// `[peer_manager::PeerManagerRequestSender::send_to]`.
    pub fn send_to(
        &mut self,
        recipient: PeerId,
        protocol: ProtocolId,
        message: TMessage,
    ) -> Result<(), NetworkError> {
        // Serialize message.
        let mdata = message.to_bytes().unwrap();
        self.inner.send_to(recipient, protocol, mdata)?;
        Ok(())
    }

    /// Send a protobuf message to a many recipients. Provides a wrapper over
    /// `[peer_manager::PeerManagerRequestSender::send_to_many]`.
    pub fn send_to_many(
        &mut self,
        recipients: impl Iterator<Item = PeerId>,
        protocol: ProtocolId,
        message: TMessage,
    ) -> Result<(), NetworkError> {
        // Serialize message.
        let mdata = message.to_bytes().unwrap();
        self.inner.send_to_many(recipients, protocol, mdata)?;
        Ok(())
    }

    /// Send a protobuf rpc request to a single recipient while handling
    /// serialization and deserialization of the request and response respectively.
    /// Assumes that the request and response both have the same message type.
    pub async fn unary_rpc(
        &mut self,
        recipient: PeerId,
        protocol: ProtocolId,
        req_msg: TMessage,
        timeout: Duration,
    ) -> Result<TMessage, RpcError> {
        // serialize request
        let req_data = req_msg.to_bytes().unwrap();
        let res_data = self
            .inner
            .unary_rpc(recipient, protocol, req_data, timeout)
            .await?;
        let res_msg = TMessage::decode(res_data.as_ref())?;
        Ok(res_msg)
    }
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    smoke_test_environment::SwarmBuilder,
    storage::{db_backup, db_restore},
    test_utils::{check_create_mint_transfer_node, swarm_utils::insert_waypoint},
    workspace_builder,
    workspace_builder::workspace_root,
};
use anyhow::anyhow;
use aptos_config::config::NodeConfig;
use aptos_temppath::TempPath;
use aptos_types::{transaction::Transaction, waypoint::Waypoint};
use forge::{get_highest_synced_version, LocalNode, Node, NodeExt, SwarmExt};
use move_core_types::language_storage::CORE_CODE_ADDRESS;
use regex::Regex;
use std::{fs, process::Command, str::FromStr, time::Duration};

fn update_node_config_restart(validator: &mut LocalNode, mut config: NodeConfig) {
    validator.stop();
    let node_path = validator.config_path();
    config.save(node_path).unwrap();
    validator.start().unwrap();
}

#[tokio::test]
/// This test verifies the flow of a genesis transaction after the chain starts.
/// 1. Test the consensus sync_only mode, every node should stop at the same version.
/// 2. Test the db-bootstrapper applying a manual genesis transaction (remove validator 0) on diemdb directly
/// 3. Test the nodes and clients resume working after updating waypoint
/// 4. Test a node lagging behind can sync to the waypoint
async fn test_genesis_transaction_flow() {
    let db_bootstrapper = workspace_builder::get_bin("db-bootstrapper");
    let aptos_cli = workspace_builder::get_bin("aptos");

    // prebuild tools.
    workspace_builder::get_bin("db-backup");
    workspace_builder::get_bin("db-restore");
    workspace_builder::get_bin("db-backup-verify");

    println!("0. pre-building finished.");

    let num_nodes = 5;
    let (mut env, cli, _) = SwarmBuilder::new_local(num_nodes)
        .with_aptos()
        .build_with_cli(0)
        .await;

    println!("1. Set sync_only = true for the last node and check it can sync to others");
    let node = env.validators_mut().nth(4).unwrap();
    let mut new_config = node.config().clone();
    new_config.consensus.sync_only = true;
    update_node_config_restart(node, new_config.clone());
    // wait for some versions
    env.wait_for_all_nodes_to_catchup_to_version(10, Duration::from_secs(10))
        .await
        .unwrap();

    println!("2. Set sync_only = true for all nodes and restart");
    for node in env.validators_mut() {
        let mut node_config = node.config().clone();
        node_config.consensus.sync_only = true;
        update_node_config_restart(node, node_config)
    }

    println!("3. delete one node's db and test they can still sync when sync_only is true for every nodes");
    let node = env.validators_mut().nth(3).unwrap();
    node.stop();
    node.clear_storage().await.unwrap();
    node.start().unwrap();

    println!("4. verify all nodes are at the same round and no progress being made");
    env.wait_for_all_nodes_to_catchup(Duration::from_secs(30))
        .await
        .unwrap();

    println!("5. kill nodes and prepare a genesis txn to remove the last validator");
    for node in env.validators_mut().take(3) {
        node.stop();
    }

    let first_validator_address = env.validators().nth(4).unwrap().config().peer_id().unwrap();

    let script = format!(
        r#"
        script {{
            use aptos_framework::stake;
            use aptos_framework::aptos_governance;
            use aptos_framework::block;

            fun main(vm_signer: &signer, framework_signer: &signer) {{
                stake::remove_validators(framework_signer, &vector[@0x{:?}]);
                block::emit_writeset_block_event(vm_signer, @0x1);
                aptos_governance::reconfigure(framework_signer);
            }}
    }}
    "#,
        first_validator_address
    );

    let temp_script_path = TempPath::new();
    let mut move_script_path = temp_script_path.path().to_path_buf();
    move_script_path.set_extension("move");

    fs::write(move_script_path.as_path(), script).unwrap();

    let genesis_blob_path = TempPath::new();
    genesis_blob_path.create_as_file().unwrap();

    Command::new(aptos_cli.as_path())
        .current_dir(workspace_root())
        .args(&vec![
            "genesis",
            "generate-admin-write-set",
            "--output-file",
            genesis_blob_path.path().to_str().unwrap(),
            "--execute-as",
            CORE_CODE_ADDRESS.clone().to_hex().as_str(),
            "--script-path",
            move_script_path.as_path().to_str().unwrap(),
            "--framework-git-rev",
            "HEAD",
            "--assume-yes",
        ])
        .output()
        .unwrap();

    let genesis_transaction = {
        let buf = fs::read(genesis_blob_path.as_ref()).unwrap();
        bcs::from_bytes::<Transaction>(&buf).unwrap()
    };

    println!("6. prepare the waypoint with the transaction");
    let waypoint_command = Command::new(db_bootstrapper.as_path())
        .current_dir(workspace_root())
        .args(&vec![
            env.validators()
                .next()
                .unwrap()
                .config()
                .storage
                .dir()
                .to_str()
                .unwrap(),
            "--genesis-txn-file",
            genesis_blob_path.path().to_str().unwrap(),
        ])
        .output()
        .unwrap();
    println!("Db bootstrapper output: {:?}", waypoint_command);
    let output = std::str::from_utf8(&waypoint_command.stdout).unwrap();
    let waypoint = parse_waypoint(output);

    println!("7. apply genesis transaction for nodes 0, 1, 2");
    for node in env.validators_mut().take(3) {
        let mut node_config = node.config().clone();
        insert_waypoint(&mut node_config, waypoint);
        node_config.execution.genesis = Some(genesis_transaction.clone());
        // reset the sync_only flag to false
        node_config.consensus.sync_only = false;
        update_node_config_restart(node, node_config);
    }

    println!("8. verify it's able to mint after the waypoint");
    env.wait_for_startup().await.unwrap();
    check_create_mint_transfer_node(&mut env, 0).await;

    let (epoch, version) = {
        let response = env
            .validators()
            .next()
            .unwrap()
            .rest_client()
            .get_ledger_information()
            .await
            .unwrap();
        (response.inner().epoch, response.inner().version)
    };

    let backup_path = db_backup(
        env.validators()
            .next()
            .unwrap()
            .config()
            .storage
            .backup_service_address
            .port(),
        epoch.checked_sub(1).unwrap(), // target epoch: most recently closed epoch
        version,                       // target version
        version as usize,              // txn batch size (version 0 is in its own batch)
        epoch.checked_sub(1).unwrap() as usize, // state snapshot interval
        &[waypoint],
    );

    println!("9. verify node 4 is out from the validator set");
    assert_eq!(
        cli.show_validator_set()
            .await
            .unwrap()
            .active_validators
            .len(),
        4
    );

    println!("10. nuke DB on node 3, and run db-restore, test if it rejoins the network okay.");
    let node = env.validators_mut().nth(3).unwrap();
    node.stop();
    let mut node_config = node.config().clone();
    node_config.consensus.sync_only = false;
    node_config.save(node.config_path()).unwrap();

    let db_dir = node.config().storage.dir();
    fs::remove_dir_all(&db_dir).unwrap();
    db_restore(backup_path.path(), db_dir.as_path(), &[waypoint]);

    node.start().unwrap();
    let client = node.rest_client();
    // wait for it to catch up
    {
        let version = get_highest_synced_version(&env.get_all_nodes_clients_with_names())
            .await
            .unwrap();
        loop {
            if let Ok(resp) = client.get_ledger_information().await {
                if resp.into_inner().version > version {
                    println!("Node 3 catches up on {}", version);
                    break;
                }
            }
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
    }

    check_create_mint_transfer_node(&mut env, 3).await;
}

fn parse_waypoint(db_bootstrapper_output: &str) -> Waypoint {
    let waypoint = Regex::new(r"Got waypoint: (\d+:\w+)")
        .unwrap()
        .captures(db_bootstrapper_output)
        .ok_or_else(|| anyhow!("Failed to parse db-bootstrapper output."));
    Waypoint::from_str(waypoint.unwrap()[1].into()).unwrap()
}
<commit_msg>[smoke tests] test_genesis_transaction_flow: add wait checks on validator restarts (#5262)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    smoke_test_environment::SwarmBuilder,
    storage::{db_backup, db_restore},
    test_utils::{check_create_mint_transfer_node, swarm_utils::insert_waypoint},
    workspace_builder,
    workspace_builder::workspace_root,
};
use anyhow::anyhow;
use aptos_config::config::NodeConfig;
use aptos_logger::prelude::*;
use aptos_temppath::TempPath;
use aptos_types::{transaction::Transaction, waypoint::Waypoint};
use forge::{get_highest_synced_version, LocalNode, Node, NodeExt, SwarmExt, Validator};
use move_core_types::language_storage::CORE_CODE_ADDRESS;
use regex::Regex;
use std::time::Instant;
use std::{fs, process::Command, str::FromStr, time::Duration};

fn update_node_config_restart(validator: &mut LocalNode, mut config: NodeConfig) {
    validator.stop();
    let node_path = validator.config_path();
    config.save(node_path).unwrap();
    validator.start().unwrap();
}

async fn wait_for_node(validator: &mut dyn Validator, expected_to_connect: usize) {
    let deadline = Instant::now().checked_add(Duration::from_secs(60)).unwrap();
    validator.wait_until_healthy(deadline).await.unwrap();
    info!("Validator restart health check passed");
    validator
        .wait_for_connectivity(expected_to_connect, deadline)
        .await
        .unwrap();
    info!("Validator restart connectivity check passed");
}

#[tokio::test]
/// This test verifies the flow of a genesis transaction after the chain starts.
/// 1. Test the consensus sync_only mode, every node should stop at the same version.
/// 2. Test the db-bootstrapper applying a manual genesis transaction (remove validator 0) on diemdb directly
/// 3. Test the nodes and clients resume working after updating waypoint
/// 4. Test a node lagging behind can sync to the waypoint
async fn test_genesis_transaction_flow() {
    let db_bootstrapper = workspace_builder::get_bin("db-bootstrapper");
    let aptos_cli = workspace_builder::get_bin("aptos");

    // prebuild tools.
    workspace_builder::get_bin("db-backup");
    workspace_builder::get_bin("db-restore");
    workspace_builder::get_bin("db-backup-verify");

    println!("0. pre-building finished.");

    let num_nodes = 5;
    let (mut env, cli, _) = SwarmBuilder::new_local(num_nodes)
        .with_aptos()
        .build_with_cli(0)
        .await;

    println!("1. Set sync_only = true for the last node and check it can sync to others");
    let node = env.validators_mut().nth(4).unwrap();
    let mut new_config = node.config().clone();
    new_config.consensus.sync_only = true;
    update_node_config_restart(node, new_config.clone());
    wait_for_node(node, num_nodes - 1).await;
    // wait for some versions
    env.wait_for_all_nodes_to_catchup_to_version(10, Duration::from_secs(10))
        .await
        .unwrap();

    println!("2. Set sync_only = true for all nodes and restart");
    for node in env.validators_mut() {
        let mut node_config = node.config().clone();
        node_config.consensus.sync_only = true;
        update_node_config_restart(node, node_config);
        wait_for_node(node, num_nodes - 1).await;
    }

    println!("3. delete one node's db and test they can still sync when sync_only is true for every nodes");
    let node = env.validators_mut().nth(3).unwrap();
    node.stop();
    node.clear_storage().await.unwrap();
    node.start().unwrap();

    println!("4. verify all nodes are at the same round and no progress being made");
    env.wait_for_all_nodes_to_catchup(Duration::from_secs(30))
        .await
        .unwrap();

    println!("5. kill nodes and prepare a genesis txn to remove the last validator");
    for node in env.validators_mut().take(3) {
        node.stop();
    }

    let first_validator_address = env.validators().nth(4).unwrap().config().peer_id().unwrap();

    let script = format!(
        r#"
        script {{
            use aptos_framework::stake;
            use aptos_framework::aptos_governance;
            use aptos_framework::block;

            fun main(vm_signer: &signer, framework_signer: &signer) {{
                stake::remove_validators(framework_signer, &vector[@0x{:?}]);
                block::emit_writeset_block_event(vm_signer, @0x1);
                aptos_governance::reconfigure(framework_signer);
            }}
    }}
    "#,
        first_validator_address
    );

    let temp_script_path = TempPath::new();
    let mut move_script_path = temp_script_path.path().to_path_buf();
    move_script_path.set_extension("move");

    fs::write(move_script_path.as_path(), script).unwrap();

    let genesis_blob_path = TempPath::new();
    genesis_blob_path.create_as_file().unwrap();

    Command::new(aptos_cli.as_path())
        .current_dir(workspace_root())
        .args(&vec![
            "genesis",
            "generate-admin-write-set",
            "--output-file",
            genesis_blob_path.path().to_str().unwrap(),
            "--execute-as",
            CORE_CODE_ADDRESS.clone().to_hex().as_str(),
            "--script-path",
            move_script_path.as_path().to_str().unwrap(),
            "--framework-git-rev",
            "HEAD",
            "--assume-yes",
        ])
        .output()
        .unwrap();

    let genesis_transaction = {
        let buf = fs::read(genesis_blob_path.as_ref()).unwrap();
        bcs::from_bytes::<Transaction>(&buf).unwrap()
    };

    println!("6. prepare the waypoint with the transaction");
    let waypoint_command = Command::new(db_bootstrapper.as_path())
        .current_dir(workspace_root())
        .args(&vec![
            env.validators()
                .next()
                .unwrap()
                .config()
                .storage
                .dir()
                .to_str()
                .unwrap(),
            "--genesis-txn-file",
            genesis_blob_path.path().to_str().unwrap(),
        ])
        .output()
        .unwrap();
    println!("Db bootstrapper output: {:?}", waypoint_command);
    let output = std::str::from_utf8(&waypoint_command.stdout).unwrap();
    let waypoint = parse_waypoint(output);

    println!("7. apply genesis transaction for nodes 0, 1, 2");
    for (expected_to_connect, node) in env.validators_mut().take(3).enumerate() {
        let mut node_config = node.config().clone();
        insert_waypoint(&mut node_config, waypoint);
        node_config.execution.genesis = Some(genesis_transaction.clone());
        // reset the sync_only flag to false
        node_config.consensus.sync_only = false;
        update_node_config_restart(node, node_config);
        wait_for_node(node, expected_to_connect).await;
    }

    println!("8. verify it's able to mint after the waypoint");
    env.wait_for_startup().await.unwrap();
    check_create_mint_transfer_node(&mut env, 0).await;

    let (epoch, version) = {
        let response = env
            .validators()
            .next()
            .unwrap()
            .rest_client()
            .get_ledger_information()
            .await
            .unwrap();
        (response.inner().epoch, response.inner().version)
    };

    let backup_path = db_backup(
        env.validators()
            .next()
            .unwrap()
            .config()
            .storage
            .backup_service_address
            .port(),
        epoch.checked_sub(1).unwrap(), // target epoch: most recently closed epoch
        version,                       // target version
        version as usize,              // txn batch size (version 0 is in its own batch)
        epoch.checked_sub(1).unwrap() as usize, // state snapshot interval
        &[waypoint],
    );

    println!("9. verify node 4 is out from the validator set");
    assert_eq!(
        cli.show_validator_set()
            .await
            .unwrap()
            .active_validators
            .len(),
        4
    );

    println!("10. nuke DB on node 3, and run db-restore, test if it rejoins the network okay.");
    let node = env.validators_mut().nth(3).unwrap();
    node.stop();
    let mut node_config = node.config().clone();
    node_config.consensus.sync_only = false;
    node_config.save(node.config_path()).unwrap();

    let db_dir = node.config().storage.dir();
    fs::remove_dir_all(&db_dir).unwrap();
    db_restore(backup_path.path(), db_dir.as_path(), &[waypoint]);

    node.start().unwrap();
    wait_for_node(node, num_nodes - 2).await;
    let client = node.rest_client();
    // wait for it to catch up
    {
        let version = get_highest_synced_version(&env.get_all_nodes_clients_with_names())
            .await
            .unwrap();
        loop {
            if let Ok(resp) = client.get_ledger_information().await {
                if resp.into_inner().version > version {
                    println!("Node 3 catches up on {}", version);
                    break;
                }
            }
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
    }

    check_create_mint_transfer_node(&mut env, 3).await;
}

fn parse_waypoint(db_bootstrapper_output: &str) -> Waypoint {
    let waypoint = Regex::new(r"Got waypoint: (\d+:\w+)")
        .unwrap()
        .captures(db_bootstrapper_output)
        .ok_or_else(|| anyhow!("Failed to parse db-bootstrapper output."));
    Waypoint::from_str(waypoint.unwrap()[1].into()).unwrap()
}
<|endoftext|>


<commit_before>use futures::{Async, Future, Poll};
use futures_cpupool::{CpuFuture, CpuPool};
use postgres_protocol::message::frontend;
use state_machine_future::RentToOwn;
use std::error::Error as StdError;
use std::io;
use std::net::{SocketAddr, ToSocketAddrs};
use std::time::{Duration, Instant};
use std::vec;
use tokio_io::io::{flush, read_exact, write_all, Flush, ReadExact, WriteAll};
use tokio_tcp::{self, TcpStream};
use tokio_timer::Delay;

#[cfg(unix)]
use tokio_uds::{self, UnixStream};

use error::{self, Error};
use params::{ConnectParams, Host};
use proto::socket::Socket;
use tls::{self, TlsConnect, TlsStream};
use {bad_response, TlsMode};

lazy_static! {
    static ref DNS_POOL: CpuPool = CpuPool::new(2);
}

#[derive(StateMachineFuture)]
pub enum Connect {
    #[state_machine_future(start)]
    #[cfg_attr(unix, state_machine_future(transitions(ResolvingDns, ConnectingUnix)))]
    #[cfg_attr(not(unix), state_machine_future(transitions(ResolvingDns)))]
    Start { params: ConnectParams, tls: TlsMode },
    #[state_machine_future(transitions(ConnectingTcp))]
    ResolvingDns {
        future: CpuFuture<vec::IntoIter<SocketAddr>, io::Error>,
        timeout: Option<Duration>,
        params: ConnectParams,
        tls: TlsMode,
    },
    #[state_machine_future(transitions(PreparingSsl))]
    ConnectingTcp {
        addrs: vec::IntoIter<SocketAddr>,
        future: tokio_tcp::ConnectFuture,
        timeout: Option<(Duration, Delay)>,
        params: ConnectParams,
        tls: TlsMode,
    },
    #[cfg(unix)]
    #[state_machine_future(transitions(PreparingSsl))]
    ConnectingUnix {
        future: tokio_uds::ConnectFuture,
        timeout: Option<Delay>,
        params: ConnectParams,
        tls: TlsMode,
    },
    #[state_machine_future(transitions(Ready, SendingSsl))]
    PreparingSsl {
        socket: Socket,
        params: ConnectParams,
        tls: TlsMode,
    },
    #[state_machine_future(transitions(FlushingSsl))]
    SendingSsl {
        future: WriteAll<Socket, Vec<u8>>,
        params: ConnectParams,
        connector: Box<TlsConnect>,
        required: bool,
    },
    #[state_machine_future(transitions(ReadingSsl))]
    FlushingSsl {
        future: Flush<Socket>,
        params: ConnectParams,
        connector: Box<TlsConnect>,
        required: bool,
    },
    #[state_machine_future(transitions(ConnectingTls, Ready))]
    ReadingSsl {
        future: ReadExact<Socket, [u8; 1]>,
        params: ConnectParams,
        connector: Box<TlsConnect>,
        required: bool,
    },
    #[state_machine_future(transitions(Ready))]
    ConnectingTls {
        future:
            Box<Future<Item = Box<TlsStream>, Error = Box<StdError + Sync + Send>> + Sync + Send>,
        params: ConnectParams,
    },
    #[state_machine_future(ready)]
    Ready(Box<TlsStream>),
    #[state_machine_future(error)]
    Failed(Error),
}

impl PollConnect for Connect {
    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
        let state = state.take();

        let timeout = state.params.connect_timeout();
        let port = state.params.port();

        match state.params.host().clone() {
            Host::Tcp(host) => transition!(ResolvingDns {
                future: DNS_POOL.spawn_fn(move || (&*host, port).to_socket_addrs()),
                params: state.params,
                tls: state.tls,
                timeout,
            }),
            #[cfg(unix)]
            Host::Unix(mut path) => {
                path.push(format!(".s.PGSQL.{}", port));
                transition!(ConnectingUnix {
                    future: UnixStream::connect(path),
                    timeout: timeout.map(|t| Delay::new(Instant::now() + t)),
                    params: state.params,
                    tls: state.tls,
                })
            }
        }
    }

    fn poll_resolving_dns<'a>(
        state: &'a mut RentToOwn<'a, ResolvingDns>,
    ) -> Poll<AfterResolvingDns, Error> {
        let mut addrs = try_ready!(state.future.poll());
        let state = state.take();

        let addr = match addrs.next() {
            Some(addr) => addr,
            None => {
                return Err(io::Error::new(io::ErrorKind::Other, "resolved to 0 addresses").into())
            }
        };

        transition!(ConnectingTcp {
            addrs,
            future: TcpStream::connect(&addr),
            timeout: state.timeout.map(|t| (t, Delay::new(Instant::now() + t))),
            params: state.params,
            tls: state.tls,
        })
    }

    fn poll_connecting_tcp<'a>(
        state: &'a mut RentToOwn<'a, ConnectingTcp>,
    ) -> Poll<AfterConnectingTcp, Error> {
        loop {
            let error = match state.future.poll() {
                Ok(Async::Ready(socket)) => {
                    let state = state.take();
                    transition!(PreparingSsl {
                        socket: Socket::Tcp(socket),
                        params: state.params,
                        tls: state.tls,
                    })
                }
                Ok(Async::NotReady) => match state.timeout {
                    Some((_, ref mut delay)) => {
                        try_ready!(
                            delay
                                .poll()
                                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                        );
                        io::Error::new(io::ErrorKind::TimedOut, "connection timed out")
                    }
                    None => return Ok(Async::NotReady),
                },
                Err(e) => e,
            };

            let addr = match state.addrs.next() {
                Some(addr) => addr,
                None => return Err(error.into()),
            };

            state.future = TcpStream::connect(&addr);
            if let Some((timeout, ref mut delay)) = state.timeout {
                delay.reset(Instant::now() + timeout);
            }
        }
    }

    #[cfg(unix)]
    fn poll_connecting_unix<'a>(
        state: &'a mut RentToOwn<'a, ConnectingUnix>,
    ) -> Poll<AfterConnectingUnix, Error> {
        match state.future.poll()? {
            Async::Ready(socket) => {
                let state = state.take();
                transition!(PreparingSsl {
                    socket: Socket::Unix(socket),
                    params: state.params,
                    tls: state.tls,
                })
            }
            Async::NotReady => match state.timeout {
                Some(ref mut delay) => {
                    try_ready!(
                        delay
                            .poll()
                            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                    );
                    Err(io::Error::new(io::ErrorKind::TimedOut, "connection timed out").into())
                }
                None => Ok(Async::NotReady),
            },
        }
    }

    fn poll_preparing_ssl<'a>(
        state: &'a mut RentToOwn<'a, PreparingSsl>,
    ) -> Poll<AfterPreparingSsl, Error> {
        let state = state.take();

        let (connector, required) = match state.tls {
            TlsMode::None => {
                transition!(Ready(Box::new(state.socket)));
            }
            TlsMode::Prefer(connector) => (connector, false),
            TlsMode::Require(connector) => (connector, true),
        };

        let mut buf = vec![];
        frontend::ssl_request(&mut buf);
        transition!(SendingSsl {
            future: write_all(state.socket, buf),
            params: state.params,
            connector,
            required,
        })
    }

    fn poll_sending_ssl<'a>(
        state: &'a mut RentToOwn<'a, SendingSsl>,
    ) -> Poll<AfterSendingSsl, Error> {
        let (stream, _) = try_ready!(state.future.poll());
        let state = state.take();
        transition!(FlushingSsl {
            future: flush(stream),
            params: state.params,
            connector: state.connector,
            required: state.required,
        })
    }

    fn poll_flushing_ssl<'a>(
        state: &'a mut RentToOwn<'a, FlushingSsl>,
    ) -> Poll<AfterFlushingSsl, Error> {
        let stream = try_ready!(state.future.poll());
        let state = state.take();
        transition!(ReadingSsl {
            future: read_exact(stream, [0]),
            params: state.params,
            connector: state.connector,
            required: state.required,
        })
    }

    fn poll_reading_ssl<'a>(
        state: &'a mut RentToOwn<'a, ReadingSsl>,
    ) -> Poll<AfterReadingSsl, Error> {
        let (stream, buf) = try_ready!(state.future.poll());
        let state = state.take();

        match buf[0] {
            b'S' => {
                let future = match state.params.host() {
                    Host::Tcp(domain) => state.connector.connect(domain, tls::Socket(stream)),
                    Host::Unix(_) => {
                        return Err(error::tls("TLS over unix sockets not supported".into()))
                    }
                };
                transition!(ConnectingTls {
                    future,
                    params: state.params,
                })
            }
            b'N' if !state.required => transition!(Ready(Box::new(stream))),
            b'N' => Err(error::tls("TLS was required but not supported".into())),
            _ => Err(bad_response()),
        }
    }

    fn poll_connecting_tls<'a>(
        state: &'a mut RentToOwn<'a, ConnectingTls>,
    ) -> Poll<AfterConnectingTls, Error> {
        let stream = try_ready!(state.future.poll().map_err(error::tls));
        transition!(Ready(stream))
    }
}

impl ConnectFuture {
    pub fn new(params: ConnectParams, tls: TlsMode) -> ConnectFuture {
        Connect::start(params, tls)
    }
}
<commit_msg>Set TCP_NODELAY for tokio-postgres<commit_after>use futures::{Async, Future, Poll};
use futures_cpupool::{CpuFuture, CpuPool};
use postgres_protocol::message::frontend;
use state_machine_future::RentToOwn;
use std::error::Error as StdError;
use std::io;
use std::net::{SocketAddr, ToSocketAddrs};
use std::time::{Duration, Instant};
use std::vec;
use tokio_io::io::{flush, read_exact, write_all, Flush, ReadExact, WriteAll};
use tokio_tcp::{self, TcpStream};
use tokio_timer::Delay;

#[cfg(unix)]
use tokio_uds::{self, UnixStream};

use error::{self, Error};
use params::{ConnectParams, Host};
use proto::socket::Socket;
use tls::{self, TlsConnect, TlsStream};
use {bad_response, TlsMode};

lazy_static! {
    static ref DNS_POOL: CpuPool = CpuPool::new(2);
}

#[derive(StateMachineFuture)]
pub enum Connect {
    #[state_machine_future(start)]
    #[cfg_attr(unix, state_machine_future(transitions(ResolvingDns, ConnectingUnix)))]
    #[cfg_attr(not(unix), state_machine_future(transitions(ResolvingDns)))]
    Start { params: ConnectParams, tls: TlsMode },
    #[state_machine_future(transitions(ConnectingTcp))]
    ResolvingDns {
        future: CpuFuture<vec::IntoIter<SocketAddr>, io::Error>,
        timeout: Option<Duration>,
        params: ConnectParams,
        tls: TlsMode,
    },
    #[state_machine_future(transitions(PreparingSsl))]
    ConnectingTcp {
        addrs: vec::IntoIter<SocketAddr>,
        future: tokio_tcp::ConnectFuture,
        timeout: Option<(Duration, Delay)>,
        params: ConnectParams,
        tls: TlsMode,
    },
    #[cfg(unix)]
    #[state_machine_future(transitions(PreparingSsl))]
    ConnectingUnix {
        future: tokio_uds::ConnectFuture,
        timeout: Option<Delay>,
        params: ConnectParams,
        tls: TlsMode,
    },
    #[state_machine_future(transitions(Ready, SendingSsl))]
    PreparingSsl {
        socket: Socket,
        params: ConnectParams,
        tls: TlsMode,
    },
    #[state_machine_future(transitions(FlushingSsl))]
    SendingSsl {
        future: WriteAll<Socket, Vec<u8>>,
        params: ConnectParams,
        connector: Box<TlsConnect>,
        required: bool,
    },
    #[state_machine_future(transitions(ReadingSsl))]
    FlushingSsl {
        future: Flush<Socket>,
        params: ConnectParams,
        connector: Box<TlsConnect>,
        required: bool,
    },
    #[state_machine_future(transitions(ConnectingTls, Ready))]
    ReadingSsl {
        future: ReadExact<Socket, [u8; 1]>,
        params: ConnectParams,
        connector: Box<TlsConnect>,
        required: bool,
    },
    #[state_machine_future(transitions(Ready))]
    ConnectingTls {
        future:
            Box<Future<Item = Box<TlsStream>, Error = Box<StdError + Sync + Send>> + Sync + Send>,
        params: ConnectParams,
    },
    #[state_machine_future(ready)]
    Ready(Box<TlsStream>),
    #[state_machine_future(error)]
    Failed(Error),
}

impl PollConnect for Connect {
    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
        let state = state.take();

        let timeout = state.params.connect_timeout();
        let port = state.params.port();

        match state.params.host().clone() {
            Host::Tcp(host) => transition!(ResolvingDns {
                future: DNS_POOL.spawn_fn(move || (&*host, port).to_socket_addrs()),
                params: state.params,
                tls: state.tls,
                timeout,
            }),
            #[cfg(unix)]
            Host::Unix(mut path) => {
                path.push(format!(".s.PGSQL.{}", port));
                transition!(ConnectingUnix {
                    future: UnixStream::connect(path),
                    timeout: timeout.map(|t| Delay::new(Instant::now() + t)),
                    params: state.params,
                    tls: state.tls,
                })
            }
        }
    }

    fn poll_resolving_dns<'a>(
        state: &'a mut RentToOwn<'a, ResolvingDns>,
    ) -> Poll<AfterResolvingDns, Error> {
        let mut addrs = try_ready!(state.future.poll());
        let state = state.take();

        let addr = match addrs.next() {
            Some(addr) => addr,
            None => {
                return Err(io::Error::new(io::ErrorKind::Other, "resolved to 0 addresses").into())
            }
        };

        transition!(ConnectingTcp {
            addrs,
            future: TcpStream::connect(&addr),
            timeout: state.timeout.map(|t| (t, Delay::new(Instant::now() + t))),
            params: state.params,
            tls: state.tls,
        })
    }

    fn poll_connecting_tcp<'a>(
        state: &'a mut RentToOwn<'a, ConnectingTcp>,
    ) -> Poll<AfterConnectingTcp, Error> {
        let socket = loop {
            let error = match state.future.poll() {
                Ok(Async::Ready(socket)) => break socket,
                Ok(Async::NotReady) => match state.timeout {
                    Some((_, ref mut delay)) => {
                        try_ready!(
                            delay
                                .poll()
                                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                        );
                        io::Error::new(io::ErrorKind::TimedOut, "connection timed out")
                    }
                    None => return Ok(Async::NotReady),
                },
                Err(e) => e,
            };

            let addr = match state.addrs.next() {
                Some(addr) => addr,
                None => return Err(error.into()),
            };

            state.future = TcpStream::connect(&addr);
            if let Some((timeout, ref mut delay)) = state.timeout {
                delay.reset(Instant::now() + timeout);
            }
        };

        // Our read/write patterns may trigger Nagle's algorithm since we're pipelining which
        // we don't want. Each individual write should be a full command we want the backend to
        // see immediately.
        socket.set_nodelay(true)?;

        let state = state.take();
        transition!(PreparingSsl {
            socket: Socket::Tcp(socket),
            params: state.params,
            tls: state.tls,
        })
    }

    #[cfg(unix)]
    fn poll_connecting_unix<'a>(
        state: &'a mut RentToOwn<'a, ConnectingUnix>,
    ) -> Poll<AfterConnectingUnix, Error> {
        match state.future.poll()? {
            Async::Ready(socket) => {
                let state = state.take();
                transition!(PreparingSsl {
                    socket: Socket::Unix(socket),
                    params: state.params,
                    tls: state.tls,
                })
            }
            Async::NotReady => match state.timeout {
                Some(ref mut delay) => {
                    try_ready!(
                        delay
                            .poll()
                            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                    );
                    Err(io::Error::new(io::ErrorKind::TimedOut, "connection timed out").into())
                }
                None => Ok(Async::NotReady),
            },
        }
    }

    fn poll_preparing_ssl<'a>(
        state: &'a mut RentToOwn<'a, PreparingSsl>,
    ) -> Poll<AfterPreparingSsl, Error> {
        let state = state.take();

        let (connector, required) = match state.tls {
            TlsMode::None => {
                transition!(Ready(Box::new(state.socket)));
            }
            TlsMode::Prefer(connector) => (connector, false),
            TlsMode::Require(connector) => (connector, true),
        };

        let mut buf = vec![];
        frontend::ssl_request(&mut buf);
        transition!(SendingSsl {
            future: write_all(state.socket, buf),
            params: state.params,
            connector,
            required,
        })
    }

    fn poll_sending_ssl<'a>(
        state: &'a mut RentToOwn<'a, SendingSsl>,
    ) -> Poll<AfterSendingSsl, Error> {
        let (stream, _) = try_ready!(state.future.poll());
        let state = state.take();
        transition!(FlushingSsl {
            future: flush(stream),
            params: state.params,
            connector: state.connector,
            required: state.required,
        })
    }

    fn poll_flushing_ssl<'a>(
        state: &'a mut RentToOwn<'a, FlushingSsl>,
    ) -> Poll<AfterFlushingSsl, Error> {
        let stream = try_ready!(state.future.poll());
        let state = state.take();
        transition!(ReadingSsl {
            future: read_exact(stream, [0]),
            params: state.params,
            connector: state.connector,
            required: state.required,
        })
    }

    fn poll_reading_ssl<'a>(
        state: &'a mut RentToOwn<'a, ReadingSsl>,
    ) -> Poll<AfterReadingSsl, Error> {
        let (stream, buf) = try_ready!(state.future.poll());
        let state = state.take();

        match buf[0] {
            b'S' => {
                let future = match state.params.host() {
                    Host::Tcp(domain) => state.connector.connect(domain, tls::Socket(stream)),
                    Host::Unix(_) => {
                        return Err(error::tls("TLS over unix sockets not supported".into()))
                    }
                };
                transition!(ConnectingTls {
                    future,
                    params: state.params,
                })
            }
            b'N' if !state.required => transition!(Ready(Box::new(stream))),
            b'N' => Err(error::tls("TLS was required but not supported".into())),
            _ => Err(bad_response()),
        }
    }

    fn poll_connecting_tls<'a>(
        state: &'a mut RentToOwn<'a, ConnectingTls>,
    ) -> Poll<AfterConnectingTls, Error> {
        let stream = try_ready!(state.future.poll().map_err(error::tls));
        transition!(Ready(stream))
    }
}

impl ConnectFuture {
    pub fn new(params: ConnectParams, tls: TlsMode) -> ConnectFuture {
        Connect::start(params, tls)
    }
}
<|endoftext|>