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

use crate::{
    constants::{INBOUND_RPC_TIMEOUT_MS, MAX_CONCURRENT_INBOUND_RPCS, MAX_FRAME_SIZE},
    peer::{DisconnectReason, Peer, PeerHandle, PeerNotification},
    protocols::{
        direct_send::Message,
        wire::{
            handshake::v1::MessagingProtocolVersion,
            messaging::v1::{
                DirectSendMsg, NetworkMessage, NetworkMessageSink, NetworkMessageStream,
            },
        },
    },
    transport::{Connection, ConnectionId, ConnectionMetadata},
    ProtocolId,
};
use bytes::Bytes;
use diem_config::network_id::NetworkContext;
use diem_network_address::NetworkAddress;
use diem_types::PeerId;
use futures::{
    future,
    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
    stream::StreamExt,
    SinkExt,
};
use memsocket::MemorySocket;
use netcore::{compat::IoCompat, transport::ConnectionOrigin};
use std::{str::FromStr, time::Duration};
use tokio::runtime::{Handle, Runtime};

static PROTOCOL: ProtocolId = ProtocolId::MempoolDirectSend;

fn build_test_peer(
    executor: Handle,
    origin: ConnectionOrigin,
) -> (
    Peer<MemorySocket>,
    PeerHandle,
    MemorySocket,
    channel::Receiver<PeerNotification>,
    channel::Receiver<PeerNotification>,
) {
    let (a, b) = MemorySocket::new_pair();
    let peer_id = PeerId::random();
    let (peer_notifs_tx, peer_notifs_rx) = channel::new_test(1);
    let (peer_rpc_notifs_tx, peer_rpc_notifs_rx) = channel::new_test(1);
    let (peer_req_tx, peer_req_rx) = channel::new_test(0);
    let connection = Connection {
        metadata: ConnectionMetadata::new(
            peer_id,
            ConnectionId::default(),
            NetworkAddress::from_str("/ip4/127.0.0.1/tcp/8081").unwrap(),
            origin,
            MessagingProtocolVersion::V1,
            [].iter().into(),
        ),
        socket: a,
    };
    let connection_metadata = connection.metadata.clone();

    let peer = Peer::new(
        NetworkContext::mock(),
        executor,
        connection,
        peer_req_rx,
        peer_notifs_tx,
        peer_rpc_notifs_tx,
        Duration::from_millis(INBOUND_RPC_TIMEOUT_MS),
        MAX_CONCURRENT_INBOUND_RPCS,
        MAX_FRAME_SIZE,
    );
    let peer_handle = PeerHandle::new(NetworkContext::mock(), connection_metadata, peer_req_tx);

    (peer, peer_handle, b, peer_notifs_rx, peer_rpc_notifs_rx)
}

fn build_test_connected_peers(
    executor: Handle,
) -> (
    (
        Peer<MemorySocket>,
        PeerHandle,
        channel::Receiver<PeerNotification>,
        channel::Receiver<PeerNotification>,
    ),
    (
        Peer<MemorySocket>,
        PeerHandle,
        channel::Receiver<PeerNotification>,
        channel::Receiver<PeerNotification>,
    ),
) {
    let (peer_a, peer_handle_a, connection_a, peer_notifs_rx_a, peer_rpc_notifs_rx_a) =
        build_test_peer(executor.clone(), ConnectionOrigin::Inbound);
    let (mut peer_b, peer_handle_b, _connection_b, peer_notifs_rx_b, peer_rpc_notifs_rx_b) =
        build_test_peer(executor, ConnectionOrigin::Outbound);
    // Make sure both peers are connected
    peer_b.connection = Some(connection_a);
    (
        (
            peer_a,
            peer_handle_a,
            peer_notifs_rx_a,
            peer_rpc_notifs_rx_a,
        ),
        (
            peer_b,
            peer_handle_b,
            peer_notifs_rx_b,
            peer_rpc_notifs_rx_b,
        ),
    )
}

fn build_network_sink_stream<'a>(
    connection: &'a mut MemorySocket,
) -> (
    NetworkMessageSink<impl AsyncWrite + 'a>,
    NetworkMessageStream<impl AsyncRead + 'a>,
) {
    let (read_half, write_half) = tokio::io::split(IoCompat::new(connection));
    let sink = NetworkMessageSink::new(IoCompat::new(write_half), MAX_FRAME_SIZE);
    let stream = NetworkMessageStream::new(IoCompat::new(read_half), MAX_FRAME_SIZE);
    (sink, stream)
}

async fn assert_peer_disconnected_event(
    peer_id: PeerId,
    reason: DisconnectReason,
    peer_notifs_rx: &mut channel::Receiver<PeerNotification>,
) {
    match peer_notifs_rx.next().await {
        Some(PeerNotification::PeerDisconnected(conn_info, actual_reason)) => {
            assert_eq!(conn_info.remote_peer_id, peer_id);
            assert_eq!(actual_reason, reason);
        }
        event => panic!("Expected a PeerDisconnected, received: {:?}", event),
    }
}

// Sending an outbound DirectSend should write it to the wire.
#[test]
fn peer_send_message() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, mut peer_handle, mut connection, _peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);
    let (mut client_sink, mut client_stream) = build_network_sink_stream(&mut connection);

    let send_msg = Message {
        protocol_id: PROTOCOL,
        mdata: Bytes::from("hello world"),
    };
    let recv_msg = NetworkMessage::DirectSendMsg(DirectSendMsg {
        protocol_id: PROTOCOL,
        priority: 0,
        raw_msg: Vec::from("hello world"),
    });

    let client = async {
        // Client should receive the direct send messages.
        for _ in 0..30 {
            let msg = client_stream.next().await.unwrap().unwrap();
            assert_eq!(msg, recv_msg);
        }
        // Client then closes the connection.
        client_sink.close().await.unwrap();
    };

    let server = async {
        // Server sends some direct send messages.
        for _ in 0..30 {
            peer_handle
                .send_direct_send(send_msg.clone())
                .await
                .unwrap();
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// Reading an inbound DirectSendMsg off the wire should notify the PeerManager of
// an inbound DirectSend.
#[test]
fn peer_recv_message() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, _peer_handle, connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let send_msg = NetworkMessage::DirectSendMsg(DirectSendMsg {
        protocol_id: PROTOCOL,
        priority: 0,
        raw_msg: Vec::from("hello world"),
    });
    let recv_msg = PeerNotification::RecvDirectSend(Message {
        protocol_id: PROTOCOL,
        mdata: Bytes::from("hello world"),
    });

    let client = async move {
        let mut connection = NetworkMessageSink::new(connection, MAX_FRAME_SIZE);
        for _ in 0..30 {
            // The client should then send the network message.
            connection.send(&send_msg).await.unwrap();
        }
        // Client then closes connection.
        connection.close().await.unwrap();
    };

    let server = async move {
        for _ in 0..30 {
            // Wait to receive notification of DirectSendMsg from Peer.
            let received = peer_notifs_rx.next().await.unwrap();
            assert_eq!(recv_msg, received);
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// Two connected Peer actors should be able to send/recv a DirectSend from each
// other and then shutdown gracefully.
#[test]
fn peers_send_message_concurrent() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (
        (peer_a, mut peer_handle_a, mut peer_notifs_rx_a, _peer_rpc_notifs_rx_a),
        (peer_b, mut peer_handle_b, mut peer_notifs_rx_b, _peer_rpc_notifs_rx_b),
    ) = build_test_connected_peers(rt.handle().clone());

    let test = async move {
        let msg_a = Message {
            protocol_id: PROTOCOL,
            mdata: Bytes::from("hello world"),
        };
        let msg_b = Message {
            protocol_id: PROTOCOL,
            mdata: Bytes::from("namaste"),
        };

        // Peer A -> msg_a -> Peer B
        peer_handle_a.send_direct_send(msg_a.clone()).await.unwrap();
        // Peer A <- msg_b <- Peer B
        peer_handle_b.send_direct_send(msg_b.clone()).await.unwrap();

        // Check that each peer received the other's message
        let notif_a = peer_notifs_rx_a.next().await;
        let notif_b = peer_notifs_rx_b.next().await;
        assert_eq!(notif_a, Some(PeerNotification::RecvDirectSend(msg_b)));
        assert_eq!(notif_b, Some(PeerNotification::RecvDirectSend(msg_a)));

        // Shut one peers and the other should shutdown due to ConnectionLost
        peer_handle_a.disconnect().await;

        // Check that we received both shutdown events
        assert_peer_disconnected_event(
            peer_handle_a.peer_id(),
            DisconnectReason::Requested,
            &mut peer_notifs_rx_a,
        )
        .await;
        assert_peer_disconnected_event(
            peer_handle_b.peer_id(),
            DisconnectReason::ConnectionLost,
            &mut peer_notifs_rx_b,
        )
        .await;
    };

    rt.block_on(future::join3(peer_a.start(), peer_b.start(), test));
}

// PeerManager can request a Peer to shutdown.
#[test]
fn peer_disconnect_request() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, mut peer_handle, _connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let test = async move {
        peer_handle.disconnect().await;
        assert_peer_disconnected_event(
            peer_handle.peer_id(),
            DisconnectReason::Requested,
            &mut peer_notifs_rx,
        )
        .await;
    };

    rt.block_on(future::join(peer.start(), test));
}

// Peer will shutdown if the underlying connection is lost.
#[test]
fn peer_disconnect_connection_lost() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, peer_handle, mut connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let test = async move {
        connection.close().await.unwrap();
        assert_peer_disconnected_event(
            peer_handle.peer_id(),
            DisconnectReason::ConnectionLost,
            &mut peer_notifs_rx,
        )
        .await;
    };
    rt.block_on(future::join(peer.start(), test));
}

#[test]
fn peer_terminates_when_request_tx_has_dropped() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, peer_handle, _connection, _peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let drop = async move {
        // Drop peer handle.
        drop(peer_handle);
    };
    rt.block_on(future::join(peer.start(), drop));
}
<commit_msg>[network] Add tests for inbound rpcs through Peer<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    constants::{INBOUND_RPC_TIMEOUT_MS, MAX_CONCURRENT_INBOUND_RPCS, MAX_FRAME_SIZE},
    peer::{DisconnectReason, Peer, PeerHandle, PeerNotification},
    protocols::{
        direct_send::Message,
        rpc::InboundRpcRequest,
        wire::{
            handshake::v1::MessagingProtocolVersion,
            messaging::v1::{
                DirectSendMsg, NetworkMessage, NetworkMessageSink, NetworkMessageStream,
                RpcRequest, RpcResponse,
            },
        },
    },
    transport::{Connection, ConnectionId, ConnectionMetadata},
    ProtocolId,
};
use bytes::Bytes;
use diem_config::network_id::NetworkContext;
use diem_network_address::NetworkAddress;
use diem_types::PeerId;
use futures::{
    channel::oneshot,
    future,
    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
    stream::{StreamExt, TryStreamExt},
    SinkExt,
};
use memsocket::MemorySocket;
use netcore::{compat::IoCompat, transport::ConnectionOrigin};
use std::{str::FromStr, time::Duration};
use tokio::runtime::{Handle, Runtime};

static PROTOCOL: ProtocolId = ProtocolId::MempoolDirectSend;

fn build_test_peer(
    executor: Handle,
    origin: ConnectionOrigin,
) -> (
    Peer<MemorySocket>,
    PeerHandle,
    MemorySocket,
    channel::Receiver<PeerNotification>,
    channel::Receiver<PeerNotification>,
) {
    let (a, b) = MemorySocket::new_pair();
    let peer_id = PeerId::random();
    let (peer_notifs_tx, peer_notifs_rx) = channel::new_test(1);
    let (peer_rpc_notifs_tx, peer_rpc_notifs_rx) = channel::new_test(1);
    let (peer_req_tx, peer_req_rx) = channel::new_test(0);
    let connection = Connection {
        metadata: ConnectionMetadata::new(
            peer_id,
            ConnectionId::default(),
            NetworkAddress::from_str("/ip4/127.0.0.1/tcp/8081").unwrap(),
            origin,
            MessagingProtocolVersion::V1,
            [].iter().into(),
        ),
        socket: a,
    };
    let connection_metadata = connection.metadata.clone();

    let peer = Peer::new(
        NetworkContext::mock(),
        executor,
        connection,
        peer_req_rx,
        peer_notifs_tx,
        peer_rpc_notifs_tx,
        Duration::from_millis(INBOUND_RPC_TIMEOUT_MS),
        MAX_CONCURRENT_INBOUND_RPCS,
        MAX_FRAME_SIZE,
    );
    let peer_handle = PeerHandle::new(NetworkContext::mock(), connection_metadata, peer_req_tx);

    (peer, peer_handle, b, peer_notifs_rx, peer_rpc_notifs_rx)
}

fn build_test_connected_peers(
    executor: Handle,
) -> (
    (
        Peer<MemorySocket>,
        PeerHandle,
        channel::Receiver<PeerNotification>,
        channel::Receiver<PeerNotification>,
    ),
    (
        Peer<MemorySocket>,
        PeerHandle,
        channel::Receiver<PeerNotification>,
        channel::Receiver<PeerNotification>,
    ),
) {
    let (peer_a, peer_handle_a, connection_a, peer_notifs_rx_a, peer_rpc_notifs_rx_a) =
        build_test_peer(executor.clone(), ConnectionOrigin::Inbound);
    let (mut peer_b, peer_handle_b, _connection_b, peer_notifs_rx_b, peer_rpc_notifs_rx_b) =
        build_test_peer(executor, ConnectionOrigin::Outbound);
    // Make sure both peers are connected
    peer_b.connection = Some(connection_a);
    (
        (
            peer_a,
            peer_handle_a,
            peer_notifs_rx_a,
            peer_rpc_notifs_rx_a,
        ),
        (
            peer_b,
            peer_handle_b,
            peer_notifs_rx_b,
            peer_rpc_notifs_rx_b,
        ),
    )
}

fn build_network_sink_stream<'a>(
    connection: &'a mut MemorySocket,
) -> (
    NetworkMessageSink<impl AsyncWrite + 'a>,
    NetworkMessageStream<impl AsyncRead + 'a>,
) {
    let (read_half, write_half) = tokio::io::split(IoCompat::new(connection));
    let sink = NetworkMessageSink::new(IoCompat::new(write_half), MAX_FRAME_SIZE);
    let stream = NetworkMessageStream::new(IoCompat::new(read_half), MAX_FRAME_SIZE);
    (sink, stream)
}

async fn assert_peer_disconnected_event(
    peer_id: PeerId,
    reason: DisconnectReason,
    peer_notifs_rx: &mut channel::Receiver<PeerNotification>,
) {
    match peer_notifs_rx.next().await {
        Some(PeerNotification::PeerDisconnected(conn_info, actual_reason)) => {
            assert_eq!(conn_info.remote_peer_id, peer_id);
            assert_eq!(actual_reason, reason);
        }
        event => panic!("Expected a PeerDisconnected, received: {:?}", event),
    }
}

// Sending an outbound DirectSend should write it to the wire.
#[test]
fn peer_send_message() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, mut peer_handle, mut connection, _peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);
    let (mut client_sink, mut client_stream) = build_network_sink_stream(&mut connection);

    let send_msg = Message {
        protocol_id: PROTOCOL,
        mdata: Bytes::from("hello world"),
    };
    let recv_msg = NetworkMessage::DirectSendMsg(DirectSendMsg {
        protocol_id: PROTOCOL,
        priority: 0,
        raw_msg: Vec::from("hello world"),
    });

    let client = async {
        // Client should receive the direct send messages.
        for _ in 0..30 {
            let msg = client_stream.next().await.unwrap().unwrap();
            assert_eq!(msg, recv_msg);
        }
        // Client then closes the connection.
        client_sink.close().await.unwrap();
    };

    let server = async {
        // Server sends some direct send messages.
        for _ in 0..30 {
            peer_handle
                .send_direct_send(send_msg.clone())
                .await
                .unwrap();
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// Reading an inbound DirectSendMsg off the wire should notify the PeerManager of
// an inbound DirectSend.
#[test]
fn peer_recv_message() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, _peer_handle, connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let send_msg = NetworkMessage::DirectSendMsg(DirectSendMsg {
        protocol_id: PROTOCOL,
        priority: 0,
        raw_msg: Vec::from("hello world"),
    });
    let recv_msg = PeerNotification::RecvDirectSend(Message {
        protocol_id: PROTOCOL,
        mdata: Bytes::from("hello world"),
    });

    let client = async move {
        let mut connection = NetworkMessageSink::new(connection, MAX_FRAME_SIZE);
        for _ in 0..30 {
            // The client should then send the network message.
            connection.send(&send_msg).await.unwrap();
        }
        // Client then closes connection.
        connection.close().await.unwrap();
    };

    let server = async move {
        for _ in 0..30 {
            // Wait to receive notification of DirectSendMsg from Peer.
            let received = peer_notifs_rx.next().await.unwrap();
            assert_eq!(recv_msg, received);
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// Two connected Peer actors should be able to send/recv a DirectSend from each
// other and then shutdown gracefully.
#[test]
fn peers_send_message_concurrent() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (
        (peer_a, mut peer_handle_a, mut peer_notifs_rx_a, _peer_rpc_notifs_rx_a),
        (peer_b, mut peer_handle_b, mut peer_notifs_rx_b, _peer_rpc_notifs_rx_b),
    ) = build_test_connected_peers(rt.handle().clone());

    let test = async move {
        let msg_a = Message {
            protocol_id: PROTOCOL,
            mdata: Bytes::from("hello world"),
        };
        let msg_b = Message {
            protocol_id: PROTOCOL,
            mdata: Bytes::from("namaste"),
        };

        // Peer A -> msg_a -> Peer B
        peer_handle_a.send_direct_send(msg_a.clone()).await.unwrap();
        // Peer A <- msg_b <- Peer B
        peer_handle_b.send_direct_send(msg_b.clone()).await.unwrap();

        // Check that each peer received the other's message
        let notif_a = peer_notifs_rx_a.next().await;
        let notif_b = peer_notifs_rx_b.next().await;
        assert_eq!(notif_a, Some(PeerNotification::RecvDirectSend(msg_b)));
        assert_eq!(notif_b, Some(PeerNotification::RecvDirectSend(msg_a)));

        // Shut one peers and the other should shutdown due to ConnectionLost
        peer_handle_a.disconnect().await;

        // Check that we received both shutdown events
        assert_peer_disconnected_event(
            peer_handle_a.peer_id(),
            DisconnectReason::Requested,
            &mut peer_notifs_rx_a,
        )
        .await;
        assert_peer_disconnected_event(
            peer_handle_b.peer_id(),
            DisconnectReason::ConnectionLost,
            &mut peer_notifs_rx_b,
        )
        .await;
    };

    rt.block_on(future::join3(peer_a.start(), peer_b.start(), test));
}

#[test]
fn peer_recv_rpc() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, _peer_handle, mut connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);
    let (mut client_sink, mut client_stream) = build_network_sink_stream(&mut connection);

    let send_msg = NetworkMessage::RpcRequest(RpcRequest {
        request_id: 123,
        protocol_id: PROTOCOL,
        priority: 0,
        raw_request: Vec::from("hello world"),
    });
    let recv_msg = PeerNotification::RecvRpc(InboundRpcRequest {
        protocol_id: PROTOCOL,
        data: Bytes::from("hello world"),
        res_tx: oneshot::channel().0,
    });
    let resp_msg = NetworkMessage::RpcResponse(RpcResponse {
        request_id: 123,
        priority: 0,
        raw_response: Vec::from("goodbye world"),
    });

    let client = async move {
        for _ in 0..30 {
            // Client should send the rpc request.
            client_sink.send(&send_msg).await.unwrap();
            // Client should then receive the expected rpc response.
            let received = client_stream.next().await.unwrap().unwrap();
            assert_eq!(received, resp_msg);
        }
        // Client then closes connection.
        client_sink.close().await.unwrap();
    };
    let server = async move {
        for _ in 0..30 {
            // Wait to receive RpcRequest from Peer.
            let received = peer_notifs_rx.next().await.unwrap();
            assert_eq!(recv_msg, received);

            // Send response to rpc.
            match received {
                PeerNotification::RecvRpc(req) => {
                    let response = Ok(Bytes::from("goodbye world"));
                    req.res_tx.send(response).unwrap()
                }
                _ => panic!("Unexpected PeerNotification: {:?}", received),
            }
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

#[test]
fn peer_recv_rpc_concurrent() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, _peer_handle, mut connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);
    let (mut client_sink, mut client_stream) = build_network_sink_stream(&mut connection);

    let send_msg = NetworkMessage::RpcRequest(RpcRequest {
        request_id: 123,
        protocol_id: PROTOCOL,
        priority: 0,
        raw_request: Vec::from("hello world"),
    });
    let recv_msg = PeerNotification::RecvRpc(InboundRpcRequest {
        protocol_id: PROTOCOL,
        data: Bytes::from("hello world"),
        res_tx: oneshot::channel().0,
    });
    let resp_msg = NetworkMessage::RpcResponse(RpcResponse {
        request_id: 123,
        priority: 0,
        raw_response: Vec::from("goodbye world"),
    });

    let client = async move {
        // The client should send many rpc requests.
        for _ in 0..30 {
            client_sink.send(&send_msg).await.unwrap();
        }

        // The client should then receive the expected rpc responses.
        for _ in 0..30 {
            let received = client_stream.next().await.unwrap().unwrap();
            assert_eq!(received, resp_msg);
        }

        // Client then closes connection.
        client_sink.close().await.unwrap();
    };
    let server = async move {
        let mut res_txs = vec![];

        // Wait to receive RpcRequests from Peer.
        for _ in 0..30 {
            let received = peer_notifs_rx.next().await.unwrap();
            assert_eq!(recv_msg, received);
            match received {
                PeerNotification::RecvRpc(req) => res_txs.push(req.res_tx),
                _ => panic!("Unexpected PeerNotification: {:?}", received),
            };
        }

        // Send all rpc responses to client.
        for res_tx in res_txs.into_iter() {
            let response = Bytes::from("goodbye world");
            res_tx.send(Ok(response)).unwrap();
        }
    };
    rt.block_on(future::join3(peer.start(), server, client));
}

// TODO(philiphayes): reenable this once mock time-service lands.
#[test]
#[ignore]
fn peer_recv_rpc_timeout() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, _peer_handle, mut connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);
    let (mut client_sink, client_stream) = build_network_sink_stream(&mut connection);

    let send_msg = NetworkMessage::RpcRequest(RpcRequest {
        request_id: 123,
        protocol_id: PROTOCOL,
        priority: 0,
        raw_request: Vec::from("hello world"),
    });
    let recv_msg = PeerNotification::RecvRpc(InboundRpcRequest {
        protocol_id: PROTOCOL,
        data: Bytes::from("hello world"),
        res_tx: oneshot::channel().0,
    });

    let test = async move {
        // Client sends the rpc request.
        client_sink.send(&send_msg).await.unwrap();

        // Server receives the rpc request from client.
        let received = peer_notifs_rx.next().await.unwrap();
        assert_eq!(received, recv_msg);

        // Pull out the request completion handle.
        let mut res_tx = match received {
            PeerNotification::RecvRpc(req) => req.res_tx,
            _ => panic!("Unexpected PeerNotification: {:?}", received),
        };

        // The rpc response channel should still be open since we haven't timed out yet.
        assert!(!res_tx.is_canceled());

        // time.advance(Duration::from_millis(INBOUND_RPC_TIMEOUT_MS).await;

        // The rpc response channel should be canceled from the timeout.
        assert!(res_tx.is_canceled());
        res_tx.cancellation().await;

        // Client then half-closes write side.
        client_sink.close().await.unwrap();

        // Client shouldn't have received any messages.
        let messages = client_stream.try_collect::<Vec<_>>().await.unwrap();
        assert_eq!(messages, vec![]);
    };
    rt.block_on(future::join(peer.start(), test));
}

#[test]
fn peer_recv_rpc_cancel() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, _peer_handle, mut connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);
    let (mut client_sink, client_stream) = build_network_sink_stream(&mut connection);

    let send_msg = NetworkMessage::RpcRequest(RpcRequest {
        request_id: 123,
        protocol_id: PROTOCOL,
        priority: 0,
        raw_request: Vec::from("hello world"),
    });
    let recv_msg = PeerNotification::RecvRpc(InboundRpcRequest {
        protocol_id: PROTOCOL,
        data: Bytes::from("hello world"),
        res_tx: oneshot::channel().0,
    });

    let test = async move {
        // Client sends the rpc request.
        client_sink.send(&send_msg).await.unwrap();

        // Server receives the rpc request from client.
        let received = peer_notifs_rx.next().await.unwrap();
        assert_eq!(received, recv_msg);

        // Pull out the request completion handle.
        let res_tx = match received {
            PeerNotification::RecvRpc(req) => req.res_tx,
            _ => panic!("Unexpected PeerNotification: {:?}", received),
        };

        // The rpc response channel should still be open since we haven't timed out yet.
        assert!(!res_tx.is_canceled());

        // Server drops the response completion handle to cancel the request.
        drop(res_tx);

        // Client then half-closes write side.
        client_sink.close().await.unwrap();

        // Client shouldn't have received any messages.
        let messages = client_stream.try_collect::<Vec<_>>().await.unwrap();
        assert_eq!(messages, vec![]);
    };
    rt.block_on(future::join(peer.start(), test));
}

// PeerManager can request a Peer to shutdown.
#[test]
fn peer_disconnect_request() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, mut peer_handle, _connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let test = async move {
        peer_handle.disconnect().await;
        assert_peer_disconnected_event(
            peer_handle.peer_id(),
            DisconnectReason::Requested,
            &mut peer_notifs_rx,
        )
        .await;
    };

    rt.block_on(future::join(peer.start(), test));
}

// Peer will shutdown if the underlying connection is lost.
#[test]
fn peer_disconnect_connection_lost() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, peer_handle, mut connection, mut peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let test = async move {
        connection.close().await.unwrap();
        assert_peer_disconnected_event(
            peer_handle.peer_id(),
            DisconnectReason::ConnectionLost,
            &mut peer_notifs_rx,
        )
        .await;
    };
    rt.block_on(future::join(peer.start(), test));
}

#[test]
fn peer_terminates_when_request_tx_has_dropped() {
    ::diem_logger::Logger::init_for_testing();
    let mut rt = Runtime::new().unwrap();
    let (peer, peer_handle, _connection, _peer_notifs_rx, _peer_rpc_notifs_rx) =
        build_test_peer(rt.handle().clone(), ConnectionOrigin::Inbound);

    let drop = async move {
        // Drop peer handle.
        drop(peer_handle);
    };
    rt.block_on(future::join(peer.start(), drop));
}
<|endoftext|>


<commit_before>use raytracer::animator::CameraKeyframe;
use raytracer::Renderer;
use scene::{Camera, Scene};
use std::num::Float;
use std::sync::{Arc, Semaphore};
use std::thread;
use vec3::Vec3;

pub struct Animator {
    pub fps: f64,
    pub animate_from: f64, // Number of frames is rounded down to nearest frame
    pub animate_to: f64,
    pub starting_frame_number: u32, // For filename
    pub renderer: Renderer
}

// TODO: Non-linear interpolation
impl Animator {
    // TODO: make this a Surface iterator so both single frame and animation
    // process flows are similar
    pub fn animate(&self, camera: Camera, shared_scene: Arc<Scene>, filename: &str) {
        let animate_start = ::time::get_time();
        let length = self.animate_to - self.animate_from;
        let total_frames = (self.fps * length).floor() as u32;

        let sema = Arc::new(Semaphore::new(1));

        for frame_number in 0..total_frames {
            let time = self.animate_from + frame_number as f64 / self.fps;
            let lerped_camera = Animator::lerp_camera(&camera, time);
            let frame_data = self.renderer.render(lerped_camera, shared_scene.clone());

            let file_frame_number = self.starting_frame_number + frame_number;
            let shared_name = format!("{}{:06}.ppm", filename, file_frame_number);

            let child_sema = sema.clone();
            sema.acquire();

            // Continue animating next frame as writing rendered frame to disk (slow) occurs
            thread::spawn(move || {
                ::util::export::to_ppm(frame_data, shared_name.as_slice());
                child_sema.release();
            });

            ::util::print_progress("*** Frame", animate_start.clone(), frame_number as usize + 1usize, total_frames as usize);
            println!("");
        }

        sema.acquire();
    }

    fn get_neighbour_keyframes(keyframes: Vec<CameraKeyframe>, time: f64)
                               -> (CameraKeyframe, CameraKeyframe, f64) {

        if keyframes.len() <= 1 {
            panic!("Not enough keyframes to interpolate: got: {} expected: >= 2", keyframes.len());
        }

        // Get the two keyframes inbetween current time
        let mut first = &keyframes[0];
        let mut second = &keyframes[1];

        for keyframe in keyframes.iter() {
            if keyframe.time <= time && time - keyframe.time >= first.time - time {
                first = keyframe;
            }

            if keyframe.time > time &&
               (keyframe.time - time < second.time - time || second.time < time) {
                second = keyframe;
            }
        }

        let keyframe_length = second.time - first.time;
        let alpha = (time - first.time) / keyframe_length;

        (first.clone(), second.clone(), alpha)
    }

    fn lerp_camera(camera: &Camera, time: f64) -> Camera {
        let keyframes = match camera.keyframes.clone() {
            Some(k) => k,
            None => panic!("Cannot lerp a camera with no keyframes!")
        };

        let (first, second, alpha) = Animator::get_neighbour_keyframes(keyframes, time);

        let lerped_position = Vec3::lerp(&first.position, &second.position, alpha);
        let lerped_look_at  = Vec3::lerp(&first.look_at, &second.look_at, alpha);
        let lerped_up       = Vec3::lerp(&first.up, &second.up, alpha);

        let mut lerped_camera = Camera::new(
            lerped_position,
            lerped_look_at,
            lerped_up,
            camera.fov_deg,
            camera.image_width,
            camera.image_height,
        );

        lerped_camera.keyframes = camera.keyframes.clone();
        lerped_camera
    }
}
<commit_msg>Fix animator camera lerp breaking on lerping at exact time of a keyframe<commit_after>use raytracer::animator::CameraKeyframe;
use raytracer::Renderer;
use scene::{Camera, Scene};
use std::num::Float;
use std::sync::{Arc, Semaphore};
use std::thread;
use vec3::Vec3;

pub struct Animator {
    pub fps: f64,
    pub animate_from: f64, // Number of frames is rounded down to nearest frame
    pub animate_to: f64,
    pub starting_frame_number: u32, // For filename
    pub renderer: Renderer
}

// TODO: Non-linear interpolation
impl Animator {
    // TODO: make this a Surface iterator so both single frame and animation
    // process flows are similar
    pub fn animate(&self, camera: Camera, shared_scene: Arc<Scene>, filename: &str) {
        let animate_start = ::time::get_time();
        let length = self.animate_to - self.animate_from;
        let total_frames = (self.fps * length).floor() as u32;

        let sema = Arc::new(Semaphore::new(1));

        for frame_number in 0..total_frames {
            let time = self.animate_from + frame_number as f64 / self.fps;
            let lerped_camera = Animator::lerp_camera(&camera, time);
            let frame_data = self.renderer.render(lerped_camera, shared_scene.clone());

            let file_frame_number = self.starting_frame_number + frame_number;
            let shared_name = format!("{}{:06}.ppm", filename, file_frame_number);

            let child_sema = sema.clone();
            sema.acquire();

            // Continue animating next frame as writing rendered frame to disk (slow) occurs
            thread::spawn(move || {
                ::util::export::to_ppm(frame_data, shared_name.as_slice());
                child_sema.release();
            });

            ::util::print_progress("*** Frame", animate_start.clone(), frame_number as usize + 1usize, total_frames as usize);
            println!("");
        }

        sema.acquire();
    }

    fn get_neighbour_keyframes(keyframes: Vec<CameraKeyframe>, time: f64)
                               -> (CameraKeyframe, CameraKeyframe, f64) {

        if keyframes.len() <= 1 {
            panic!("Not enough keyframes to interpolate: got: {} expected: >= 2", keyframes.len());
        }

        // Get the two keyframes inbetween current time
        let mut first = &keyframes[0];
        let mut second = &keyframes[1];

        for keyframe in keyframes.iter() {
            if keyframe.time <= time && time - keyframe.time >= first.time - time {
                first = keyframe;
            }

            if keyframe.time > time &&
               (keyframe.time - time < second.time - time || second.time < time) {
                second = keyframe;
            }
        }

        let keyframe_length = second.time - first.time;

        let alpha = if keyframe_length == 0.0 {
            0.0
        } else {
            (time - first.time) / keyframe_length
        };

        (first.clone(), second.clone(), alpha)
    }

    fn lerp_camera(camera: &Camera, time: f64) -> Camera {
        let keyframes = match camera.keyframes.clone() {
            Some(k) => k,
            None => panic!("Cannot lerp a camera with no keyframes!")
        };

        let (first, second, alpha) = Animator::get_neighbour_keyframes(keyframes, time);

        let lerped_position = Vec3::lerp(&first.position, &second.position, alpha);
        let lerped_look_at  = Vec3::lerp(&first.look_at, &second.look_at, alpha);
        let lerped_up       = Vec3::lerp(&first.up, &second.up, alpha);

        let mut lerped_camera = Camera::new(
            lerped_position,
            lerped_look_at,
            lerped_up,
            camera.fov_deg,
            camera.image_width,
            camera.image_height,
        );

        lerped_camera.keyframes = camera.keyframes.clone();
        lerped_camera
    }
}

#[test]
fn test_lerp_camera_position() {
    // Camera rotates 180 degrees
    let camera = Camera::new_with_keyframes(
        Vec3 { x: -1.0, y: -1.0, z: -1.0 },
        Vec3 { x: 0.0, y: 1.0, z: 0.0 },
        Vec3 { x: 0.0, y: 1.0, z: 0.0 },
        45.0,
        10,
        10,
        vec![
            CameraKeyframe {
                time: 5.0,
                position: Vec3 { x: 0.0, y: 0.0, z: 0.0 },
                look_at: Vec3 { x: 0.0, y: 1.0, z: 0.0 },
                up: Vec3 { x: 0.0, y: 1.0, z: 0.0 }
            },
            CameraKeyframe {
                time: 10.0,
                position: Vec3 { x: 10.0, y: 0.0, z: 0.0 },
                look_at: Vec3 { x: 0.0, y: 1.0, z: 0.0 },
                up: Vec3 { x: 0.0, y: 1.0, z: 0.0 }
            },
        ]
    );

    let expected_position_0 = Vec3 { x: -1.0, y: -1.0, z: -1.0 };
    assert_eq!(Animator::lerp_camera(&camera, 0.0).position, expected_position_0);

    let expected_position_5 = Vec3 { x: 0.0, y: 0.0, z: 0.0 };
    assert_eq!(Animator::lerp_camera(&camera, 5.0).position, expected_position_5);

    let expected_position_7_5 = Vec3 { x: 5.0, y: 0.0, z: 0.0 };
    assert_eq!(Animator::lerp_camera(&camera, 7.5).position, expected_position_7_5);

    let expected_position_10 = Vec3 { x: 10.0, y: 0.0, z: 0.0 };
    assert_eq!(Animator::lerp_camera(&camera, 10.0).position, expected_position_10);
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::error;
use std::ffi::{CString, CStr};
use std::fmt;
use std::ptr;
use std::str;

use Error;
use OomError;
use VulkanObject;
use VulkanPointers;
use instance::PhysicalDevice;
use instance::loader;
use instance::loader::LoadingError;
use vk;
use check_errors;

macro_rules! extensions {
    ($sname:ident, $($ext:ident => $s:expr,)*) => (
        /// List of extensions that are enabled or available.
        #[derive(Copy, Clone, PartialEq, Eq)]
        #[allow(missing_docs)]
        pub struct $sname {
            $(
                pub $ext: bool,
            )*

            /// This field ensures that an instance of this `Extensions` struct
            /// can only be created through Vulkano functions and the update
            /// syntax. This way, extensions can be added to Vulkano without
            /// breaking existing code.
            pub _unbuildable: Unbuildable,
        }

        impl $sname {
            /// Returns an `Extensions` object with all members set to `false`.
            #[inline]
            pub fn none() -> $sname {
                $sname {
                    $($ext: false,)*
                    _unbuildable: Unbuildable(())
                }
            }

            /// Builds a Vec containing the list of extensions.
            pub fn build_extensions_list(&self) -> Vec<CString> {
                let mut data = Vec::new();
                $(if self.$ext { data.push(CString::new(&$s[..]).unwrap()); })*
                data
            }

            /// Returns the intersection of this list and another list.
            #[inline]
            pub fn intersection(&self, other: &$sname) -> $sname {
                $sname {
                    $(
                        $ext: self.$ext && other.$ext,
                    )*
                    _unbuildable: Unbuildable(())
                }
            }
        }

        impl fmt::Debug for $sname {
            #[allow(unused_assignments)]
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                try!(write!(f, "["));

                let mut first = true;

                $(
                    if self.$ext {
                        if !first { try!(write!(f, ", ")); }
                        else { first = false; }
                        try!(f.write_str(str::from_utf8($s).unwrap()));
                    }
                )*

                write!(f, "]")
            }
        }
    );
}

macro_rules! instance_extensions {
    ($sname:ident, $($ext:ident => $s:expr,)*) => (
        extensions! {
            $sname,
            $( $ext => $s,)*
        }

        impl $sname {
            /// See the docs of supported_by_core().
            pub fn supported_by_core_raw() -> Result<$sname, SupportedExtensionsError> {
                let entry_points = try!(loader::entry_points());

                let properties: Vec<vk::ExtensionProperties> = unsafe {
                    let mut num = 0;
                    try!(check_errors(entry_points.EnumerateInstanceExtensionProperties(
                        ptr::null(), &mut num, ptr::null_mut())));

                    let mut properties = Vec::with_capacity(num as usize);
                    try!(check_errors(entry_points.EnumerateInstanceExtensionProperties(
                        ptr::null(), &mut num, properties.as_mut_ptr())));
                    properties.set_len(num as usize);
                    properties
                };

                let mut extensions = $sname::none();
                for property in properties {
                    let name = unsafe { CStr::from_ptr(property.extensionName.as_ptr()) };
                    $(
                        // TODO: Check specVersion?
                        if name.to_bytes() == &$s[..] {
                            extensions.$ext = true;
                        }
                    )*
                }

                Ok(extensions)
            }

            /// Returns an `Extensions` object with extensions supported by the core driver.
            pub fn supported_by_core() -> Result<$sname, LoadingError> {
                match $sname::supported_by_core_raw() {
                    Ok(l) => Ok(l),
                    Err(SupportedExtensionsError::LoadingError(e)) => Err(e),
                    Err(SupportedExtensionsError::OomError(e)) => panic!("{:?}", e),
                }
            }
        }
    );
}

macro_rules! device_extensions {
    ($sname:ident, $($ext:ident => $s:expr,)*) => (
        extensions! {
            $sname,
            $( $ext => $s,)*
        }

        impl $sname {
            /// See the docs of supported_by_device().
            pub fn supported_by_device_raw(physical_device: &PhysicalDevice) -> Result<$sname, SupportedExtensionsError> {
                let vk = physical_device.instance().pointers();

                let properties: Vec<vk::ExtensionProperties> = unsafe {
                    let mut num = 0;
                    try!(check_errors(vk.EnumerateDeviceExtensionProperties(
                        physical_device.internal_object(), ptr::null(), &mut num, ptr::null_mut())));

                    let mut properties = Vec::with_capacity(num as usize);
                    try!(check_errors(vk.EnumerateDeviceExtensionProperties(
                        physical_device.internal_object(), ptr::null(), &mut num, properties.as_mut_ptr())));
                    properties.set_len(num as usize);
                    properties
                };

                let mut extensions = $sname::none();
                for property in properties {
                    let name = unsafe { CStr::from_ptr(property.extensionName.as_ptr()) };
                    $(
                        // TODO: Check specVersion?
                        if name.to_bytes() == &$s[..] {
                            extensions.$ext = true;
                        }
                    )*
                }

                Ok(extensions)
            }

            /// Returns an `Extensions` object with extensions supported by the `PhysicalDevice`.
            pub fn supported_by_device(physical_device: &PhysicalDevice) -> $sname {
                match $sname::supported_by_device_raw(physical_device) {
                    Ok(l) => l,
                    Err(SupportedExtensionsError::LoadingError(e)) => unreachable!(),
                    Err(SupportedExtensionsError::OomError(e)) => panic!("{:?}", e),
                }
            }
        }
    );
}

instance_extensions! {
    InstanceExtensions,
    khr_surface => b"VK_KHR_surface",
    khr_display => b"VK_KHR_display",
    khr_xlib_surface => b"VK_KHR_xlib_surface",
    khr_xcb_surface => b"VK_KHR_xcb_surface",
    khr_wayland_surface => b"VK_KHR_wayland_surface",
    khr_mir_surface => b"VK_KHR_mir_surface",
    khr_android_surface => b"VK_KHR_android_surface",
    khr_win32_surface => b"VK_KHR_win32_surface",
    ext_debug_report => b"VK_EXT_debug_report",
    nn_vi_surface => b"VK_NN_vi_surface",
    ext_swapchain_colorspace => b"VK_EXT_swapchain_colorspace",
}

device_extensions! {
    DeviceExtensions,
    khr_swapchain => b"VK_KHR_swapchain",
    khr_display_swapchain => b"VK_KHR_display_swapchain",
    khr_sampler_mirror_clamp_to_edge => b"VK_KHR_sampler_mirror_clamp_to_edge",
    khr_maintenance1 => b"VK_KHR_maintenance1",
}

/// Error that can happen when loading the list of layers.
#[derive(Clone, Debug)]
pub enum SupportedExtensionsError {
    /// Failed to load the Vulkan shared library.
    LoadingError(LoadingError),
    /// Not enough memory.
    OomError(OomError),
}

impl error::Error for SupportedExtensionsError {
    #[inline]
    fn description(&self) -> &str {
        match *self {
            SupportedExtensionsError::LoadingError(_) => "failed to load the Vulkan shared library",
            SupportedExtensionsError::OomError(_) => "not enough memory available",
        }
    }

    #[inline]
    fn cause(&self) -> Option<&error::Error> {
        match *self {
            SupportedExtensionsError::LoadingError(ref err) => Some(err),
            SupportedExtensionsError::OomError(ref err) => Some(err),
        }
    }
}

impl fmt::Display for SupportedExtensionsError {
    #[inline]
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl From<OomError> for SupportedExtensionsError {
    #[inline]
    fn from(err: OomError) -> SupportedExtensionsError {
        SupportedExtensionsError::OomError(err)
    }
}

impl From<LoadingError> for SupportedExtensionsError {
    #[inline]
    fn from(err: LoadingError) -> SupportedExtensionsError {
        SupportedExtensionsError::LoadingError(err)
    }
}

impl From<Error> for SupportedExtensionsError {
    #[inline]
    fn from(err: Error) -> SupportedExtensionsError {
        match err {
            err @ Error::OutOfHostMemory => {
                SupportedExtensionsError::OomError(OomError::from(err))
            },
            err @ Error::OutOfDeviceMemory => {
                SupportedExtensionsError::OomError(OomError::from(err))
            },
            _ => panic!("unexpected error: {:?}", err)
        }
    }
}

/// This helper type can only be instantiated inside this module.
/// See `*Extensions::_unbuildable`.
#[doc(hidden)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct Unbuildable(());

#[cfg(test)]
mod tests {
    use instance::InstanceExtensions;
    use instance::DeviceExtensions;

    #[test]
    fn empty_extensions() {
        let i = InstanceExtensions::none().build_extensions_list();
        assert!(i.is_empty());

        let d = DeviceExtensions::none().build_extensions_list();
        assert!(d.is_empty());
    }
}
<commit_msg>Add khr_get_physical_device_properties_2 to list of instance extensions<commit_after>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::error;
use std::ffi::{CString, CStr};
use std::fmt;
use std::ptr;
use std::str;

use Error;
use OomError;
use VulkanObject;
use VulkanPointers;
use instance::PhysicalDevice;
use instance::loader;
use instance::loader::LoadingError;
use vk;
use check_errors;

macro_rules! extensions {
    ($sname:ident, $($ext:ident => $s:expr,)*) => (
        /// List of extensions that are enabled or available.
        #[derive(Copy, Clone, PartialEq, Eq)]
        #[allow(missing_docs)]
        pub struct $sname {
            $(
                pub $ext: bool,
            )*

            /// This field ensures that an instance of this `Extensions` struct
            /// can only be created through Vulkano functions and the update
            /// syntax. This way, extensions can be added to Vulkano without
            /// breaking existing code.
            pub _unbuildable: Unbuildable,
        }

        impl $sname {
            /// Returns an `Extensions` object with all members set to `false`.
            #[inline]
            pub fn none() -> $sname {
                $sname {
                    $($ext: false,)*
                    _unbuildable: Unbuildable(())
                }
            }

            /// Builds a Vec containing the list of extensions.
            pub fn build_extensions_list(&self) -> Vec<CString> {
                let mut data = Vec::new();
                $(if self.$ext { data.push(CString::new(&$s[..]).unwrap()); })*
                data
            }

            /// Returns the intersection of this list and another list.
            #[inline]
            pub fn intersection(&self, other: &$sname) -> $sname {
                $sname {
                    $(
                        $ext: self.$ext && other.$ext,
                    )*
                    _unbuildable: Unbuildable(())
                }
            }
        }

        impl fmt::Debug for $sname {
            #[allow(unused_assignments)]
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                try!(write!(f, "["));

                let mut first = true;

                $(
                    if self.$ext {
                        if !first { try!(write!(f, ", ")); }
                        else { first = false; }
                        try!(f.write_str(str::from_utf8($s).unwrap()));
                    }
                )*

                write!(f, "]")
            }
        }
    );
}

macro_rules! instance_extensions {
    ($sname:ident, $($ext:ident => $s:expr,)*) => (
        extensions! {
            $sname,
            $( $ext => $s,)*
        }

        impl $sname {
            /// See the docs of supported_by_core().
            pub fn supported_by_core_raw() -> Result<$sname, SupportedExtensionsError> {
                let entry_points = try!(loader::entry_points());

                let properties: Vec<vk::ExtensionProperties> = unsafe {
                    let mut num = 0;
                    try!(check_errors(entry_points.EnumerateInstanceExtensionProperties(
                        ptr::null(), &mut num, ptr::null_mut())));

                    let mut properties = Vec::with_capacity(num as usize);
                    try!(check_errors(entry_points.EnumerateInstanceExtensionProperties(
                        ptr::null(), &mut num, properties.as_mut_ptr())));
                    properties.set_len(num as usize);
                    properties
                };

                let mut extensions = $sname::none();
                for property in properties {
                    let name = unsafe { CStr::from_ptr(property.extensionName.as_ptr()) };
                    $(
                        // TODO: Check specVersion?
                        if name.to_bytes() == &$s[..] {
                            extensions.$ext = true;
                        }
                    )*
                }

                Ok(extensions)
            }

            /// Returns an `Extensions` object with extensions supported by the core driver.
            pub fn supported_by_core() -> Result<$sname, LoadingError> {
                match $sname::supported_by_core_raw() {
                    Ok(l) => Ok(l),
                    Err(SupportedExtensionsError::LoadingError(e)) => Err(e),
                    Err(SupportedExtensionsError::OomError(e)) => panic!("{:?}", e),
                }
            }
        }
    );
}

macro_rules! device_extensions {
    ($sname:ident, $($ext:ident => $s:expr,)*) => (
        extensions! {
            $sname,
            $( $ext => $s,)*
        }

        impl $sname {
            /// See the docs of supported_by_device().
            pub fn supported_by_device_raw(physical_device: &PhysicalDevice) -> Result<$sname, SupportedExtensionsError> {
                let vk = physical_device.instance().pointers();

                let properties: Vec<vk::ExtensionProperties> = unsafe {
                    let mut num = 0;
                    try!(check_errors(vk.EnumerateDeviceExtensionProperties(
                        physical_device.internal_object(), ptr::null(), &mut num, ptr::null_mut())));

                    let mut properties = Vec::with_capacity(num as usize);
                    try!(check_errors(vk.EnumerateDeviceExtensionProperties(
                        physical_device.internal_object(), ptr::null(), &mut num, properties.as_mut_ptr())));
                    properties.set_len(num as usize);
                    properties
                };

                let mut extensions = $sname::none();
                for property in properties {
                    let name = unsafe { CStr::from_ptr(property.extensionName.as_ptr()) };
                    $(
                        // TODO: Check specVersion?
                        if name.to_bytes() == &$s[..] {
                            extensions.$ext = true;
                        }
                    )*
                }

                Ok(extensions)
            }

            /// Returns an `Extensions` object with extensions supported by the `PhysicalDevice`.
            pub fn supported_by_device(physical_device: &PhysicalDevice) -> $sname {
                match $sname::supported_by_device_raw(physical_device) {
                    Ok(l) => l,
                    Err(SupportedExtensionsError::LoadingError(e)) => unreachable!(),
                    Err(SupportedExtensionsError::OomError(e)) => panic!("{:?}", e),
                }
            }
        }
    );
}

instance_extensions! {
    InstanceExtensions,
    khr_surface => b"VK_KHR_surface",
    khr_display => b"VK_KHR_display",
    khr_xlib_surface => b"VK_KHR_xlib_surface",
    khr_xcb_surface => b"VK_KHR_xcb_surface",
    khr_wayland_surface => b"VK_KHR_wayland_surface",
    khr_mir_surface => b"VK_KHR_mir_surface",
    khr_android_surface => b"VK_KHR_android_surface",
    khr_win32_surface => b"VK_KHR_win32_surface",
    ext_debug_report => b"VK_EXT_debug_report",
    nn_vi_surface => b"VK_NN_vi_surface",
    ext_swapchain_colorspace => b"VK_EXT_swapchain_colorspace",
    khr_get_physical_device_properties2 => b"VK_KHR_get_physical_device_properties2",
}

device_extensions! {
    DeviceExtensions,
    khr_swapchain => b"VK_KHR_swapchain",
    khr_display_swapchain => b"VK_KHR_display_swapchain",
    khr_sampler_mirror_clamp_to_edge => b"VK_KHR_sampler_mirror_clamp_to_edge",
    khr_maintenance1 => b"VK_KHR_maintenance1",
}

/// Error that can happen when loading the list of layers.
#[derive(Clone, Debug)]
pub enum SupportedExtensionsError {
    /// Failed to load the Vulkan shared library.
    LoadingError(LoadingError),
    /// Not enough memory.
    OomError(OomError),
}

impl error::Error for SupportedExtensionsError {
    #[inline]
    fn description(&self) -> &str {
        match *self {
            SupportedExtensionsError::LoadingError(_) => "failed to load the Vulkan shared library",
            SupportedExtensionsError::OomError(_) => "not enough memory available",
        }
    }

    #[inline]
    fn cause(&self) -> Option<&error::Error> {
        match *self {
            SupportedExtensionsError::LoadingError(ref err) => Some(err),
            SupportedExtensionsError::OomError(ref err) => Some(err),
        }
    }
}

impl fmt::Display for SupportedExtensionsError {
    #[inline]
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl From<OomError> for SupportedExtensionsError {
    #[inline]
    fn from(err: OomError) -> SupportedExtensionsError {
        SupportedExtensionsError::OomError(err)
    }
}

impl From<LoadingError> for SupportedExtensionsError {
    #[inline]
    fn from(err: LoadingError) -> SupportedExtensionsError {
        SupportedExtensionsError::LoadingError(err)
    }
}

impl From<Error> for SupportedExtensionsError {
    #[inline]
    fn from(err: Error) -> SupportedExtensionsError {
        match err {
            err @ Error::OutOfHostMemory => {
                SupportedExtensionsError::OomError(OomError::from(err))
            },
            err @ Error::OutOfDeviceMemory => {
                SupportedExtensionsError::OomError(OomError::from(err))
            },
            _ => panic!("unexpected error: {:?}", err)
        }
    }
}

/// This helper type can only be instantiated inside this module.
/// See `*Extensions::_unbuildable`.
#[doc(hidden)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct Unbuildable(());

#[cfg(test)]
mod tests {
    use instance::InstanceExtensions;
    use instance::DeviceExtensions;

    #[test]
    fn empty_extensions() {
        let i = InstanceExtensions::none().build_extensions_list();
        assert!(i.is_empty());

        let d = DeviceExtensions::none().build_extensions_list();
        assert!(d.is_empty());
    }
}
<|endoftext|>


<commit_before>// (C) 2015 Viktor Dahl <pazaconyoman@gmail.com>
// (C) 2015 Michael Howell <michael@notriddle.com>
// This file is licensed under the same terms as Rust itself.

#![feature(test)]

extern crate quickersort;
extern crate test;
extern crate rand;

use quickersort::{heapsort, insertion_sort, sort};
use rand::{weak_rng, Rng};
use std::mem;
use test::Bencher;

type BigSortable = (u64,u64,u64,u64);

macro_rules! bench_random(
    ($name: ident, $sortfun: ident, $typ: ty, $n: expr) => (
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut rng = weak_rng();
            b.iter(|| {
                let mut v = rng.gen_iter::<$typ>().take($n).collect::<Vec<$typ>>();
                $sortfun(&mut v[..]);
            });
            b.bytes = $n * mem::size_of::<$typ>() as u64;
        }
    )
);

bench_random!(sort_tiny_random_small, sort, u8, 5);
bench_random!(sort_tiny_random_medium, sort, u8, 100);
bench_random!(sort_tiny_random_large, sort, u8, 10_000);

bench_random!(sort_random_small, sort, u64, 5);
bench_random!(sort_random_medium, sort, u64, 100);
bench_random!(sort_random_large, sort, u64, 10_000);

bench_random!(sort_big_random_small, sort, BigSortable, 5);
bench_random!(sort_big_random_medium, sort, BigSortable, 100);
bench_random!(sort_big_random_large, sort, BigSortable, 10_000);

#[bench]
fn sort_sorted(b: &mut Bencher) {
    let mut v: Vec<_> = (0 .. 10000isize).collect();
    b.iter(|| {
        sort(&mut v[..]);
    });
    b.bytes = (v.len() * mem::size_of_val(&v[0])) as u64;
}

#[bench]
fn sort_big_sorted(b: &mut Bencher) {
    let mut v: Vec<_> = (0 .. 10000usize).map(|i| (i, i, i, i)).collect();
    b.iter(|| {
        sort(&mut v[..]);
    });
    b.bytes = (v.len() * mem::size_of_val(&v[0])) as u64;
}

#[bench]
fn sort_few_unique(b: &mut Bencher) {
    let mut v = Vec::new();
    for i in (0u32 .. 10) {
        for _ in (0usize .. 100) {
            v.push(i);
        }
    }
    let mut rng = weak_rng();
    b.iter(||{
        rng.shuffle(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = (v.len() * mem::size_of_val(&v[0])) as u64;
}

#[bench]
fn sort_equals(b: &mut Bencher) {
    let mut v = vec![1u64; 1000];
    b.iter(|| {
        sort(&mut v[..]);
    });
    b.bytes = (v.len() * mem::size_of_val(&v[0])) as u64;
}

#[bench]
fn sort_huge(b: &mut Bencher) {
    let mut rng = weak_rng();
    let n = 100_000;
    let mut v = rng.gen_iter::<i64>().take(n).collect::<Vec<i64>>();
    b.iter(|| {
        rng.shuffle(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = (n * mem::size_of::<i64>()) as u64;
}

#[bench]
fn sort_enormous(b: &mut Bencher) {
    let mut rng = weak_rng();
    let n = 1_00_000;
    let mut v = rng.gen_iter::<i64>().take(n).collect::<Vec<i64>>();
    b.iter(|| {
        rng.shuffle(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = (n * mem::size_of::<i64>()) as u64;
}

#[bench]
fn sort_partially_sorted(b: &mut Bencher) {
    fn partially_sort<T: Ord+::std::fmt::Display>(v: &mut [T]) {
        let s = v.len() / 100;
        if s == 0 { return; }
        let mut sorted = true;
        for c in v.chunks_mut(s) {
            if sorted { sort(&mut c[..]); }
            sorted = !sorted;
        }
    }
    let mut rng = weak_rng();
    let n = 10_000;
    let mut v = rng.gen_iter::<i64>().take(n).collect::<Vec<i64>>();
    b.iter(|| {
        rng.shuffle(&mut v[..]);
        partially_sort(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = (n * mem::size_of::<i64>()) as u64;
}

#[bench]
fn sort_random_enormous_heapsort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(1_000_000).collect::<Vec<u64>>();
        heapsort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 1_000_000 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_random_huge_heapsort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(100_000).collect::<Vec<u64>>();
        heapsort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 100_000 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_random_large_heapsort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(10_000).collect::<Vec<u64>>();
        heapsort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 10_000 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_random_medium_insertion_sort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(100).collect::<Vec<u64>>();
        insertion_sort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 100 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_random_medium_heapsort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(100).collect::<Vec<u64>>();
        heapsort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 100 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_strings(b: &mut Bencher) {
    let mut rng = weak_rng();
    let n = 10_000usize;
    let mut v = Vec::with_capacity(n);
    let mut bytes = 0;
    for _ in (0 .. n) {
        let len = rng.gen_range(0, 60);
        bytes += len;
        let mut s = String::with_capacity(len);
        if len == 0 {
            v.push(s);
            continue;
        }
        for _ in (0 .. len) {
            s.push(rng.gen_range(b'a', b'z') as char);
        }
        v.push(s);
    }

    b.iter(|| {
        rng.shuffle(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = bytes as u64;
}
<commit_msg>Fix compiler warnings and improve formatting<commit_after>// (C) 2015 Viktor Dahl <pazaconyoman@gmail.com>
// (C) 2015 Michael Howell <michael@notriddle.com>
// This file is licensed under the same terms as Rust itself.

#![feature(test)]

extern crate quickersort;
extern crate test;
extern crate rand;

use quickersort::{heapsort, insertion_sort, sort};
use rand::{weak_rng, Rng};
use std::mem;
use test::Bencher;

type BigSortable = (u64, u64, u64, u64);

macro_rules! bench_random(
    ($name: ident, $sortfun: ident, $typ: ty, $n: expr) => (
        #[bench]
        fn $name(b: &mut Bencher) {
            let mut rng = weak_rng();
            b.iter(|| {
                let mut v = rng.gen_iter::<$typ>().take($n).collect::<Vec<$typ>>();
                $sortfun(&mut v[..]);
            });
            b.bytes = $n * mem::size_of::<$typ>() as u64;
        }
    )
);

bench_random!(sort_tiny_random_small, sort, u8, 5);
bench_random!(sort_tiny_random_medium, sort, u8, 100);
bench_random!(sort_tiny_random_large, sort, u8, 10_000);

bench_random!(sort_random_small, sort, u64, 5);
bench_random!(sort_random_medium, sort, u64, 100);
bench_random!(sort_random_large, sort, u64, 10_000);

bench_random!(sort_big_random_small, sort, BigSortable, 5);
bench_random!(sort_big_random_medium, sort, BigSortable, 100);
bench_random!(sort_big_random_large, sort, BigSortable, 10_000);

#[bench]
fn sort_sorted(b: &mut Bencher) {
    let mut v: Vec<_> = (0 .. 10000isize).collect();
    b.iter(|| {
        sort(&mut v[..]);
    });
    b.bytes = (v.len() * mem::size_of_val(&v[0])) as u64;
}

#[bench]
fn sort_big_sorted(b: &mut Bencher) {
    let mut v: Vec<_> = (0 .. 10000usize).map(|i| (i, i, i, i)).collect();
    b.iter(|| {
        sort(&mut v[..]);
    });
    b.bytes = (v.len() * mem::size_of_val(&v[0])) as u64;
}

#[bench]
fn sort_few_unique(b: &mut Bencher) {
    let mut v = Vec::new();
    for i in 0u32 .. 10 {
        for _ in 0usize .. 100 {
            v.push(i);
        }
    }
    let mut rng = weak_rng();
    b.iter(||{
        rng.shuffle(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = (v.len() * mem::size_of_val(&v[0])) as u64;
}

#[bench]
fn sort_equals(b: &mut Bencher) {
    let mut v = vec![1u64; 1000];
    b.iter(|| {
        sort(&mut v[..]);
    });
    b.bytes = (v.len() * mem::size_of_val(&v[0])) as u64;
}

#[bench]
fn sort_huge(b: &mut Bencher) {
    let mut rng = weak_rng();
    let n = 100_000;
    let mut v = rng.gen_iter::<i64>().take(n).collect::<Vec<i64>>();
    b.iter(|| {
        rng.shuffle(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = (n * mem::size_of::<i64>()) as u64;
}

#[bench]
fn sort_enormous(b: &mut Bencher) {
    let mut rng = weak_rng();
    let n = 1_00_000;
    let mut v = rng.gen_iter::<i64>().take(n).collect::<Vec<i64>>();
    b.iter(|| {
        rng.shuffle(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = (n * mem::size_of::<i64>()) as u64;
}

#[bench]
fn sort_partially_sorted(b: &mut Bencher) {
    fn partially_sort<T: Ord+::std::fmt::Display>(v: &mut [T]) {
        let s = v.len() / 100;
        if s == 0 { return; }
        let mut sorted = true;
        for c in v.chunks_mut(s) {
            if sorted { sort(&mut c[..]); }
            sorted = !sorted;
        }
    }
    let mut rng = weak_rng();
    let n = 10_000;
    let mut v = rng.gen_iter::<i64>().take(n).collect::<Vec<i64>>();
    b.iter(|| {
        rng.shuffle(&mut v[..]);
        partially_sort(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = (n * mem::size_of::<i64>()) as u64;
}

#[bench]
fn sort_random_enormous_heapsort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(1_000_000).collect::<Vec<u64>>();
        heapsort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 1_000_000 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_random_huge_heapsort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(100_000).collect::<Vec<u64>>();
        heapsort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 100_000 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_random_large_heapsort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(10_000).collect::<Vec<u64>>();
        heapsort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 10_000 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_random_medium_insertion_sort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(100).collect::<Vec<u64>>();
        insertion_sort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 100 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_random_medium_heapsort(b: &mut Bencher) {
    let mut rng = weak_rng();
    b.iter(|| {
        let mut v = rng.gen_iter::<u64>().take(100).collect::<Vec<u64>>();
        heapsort(&mut v[..], &|a, b| a.cmp(b));
    });
    b.bytes = 100 * mem::size_of::<u64>() as u64;
}

#[bench]
fn sort_strings(b: &mut Bencher) {
    let mut rng = weak_rng();
    let n = 10_000usize;
    let mut v = Vec::with_capacity(n);
    let mut bytes = 0;
    for _ in 0 .. n {
        let len = rng.gen_range(0, 60);
        bytes += len;
        let mut s = String::with_capacity(len);
        if len == 0 {
            v.push(s);
            continue;
        }
        for _ in 0 .. len {
            s.push(rng.gen_range(b'a', b'z') as char);
        }
        v.push(s);
    }

    b.iter(|| {
        rng.shuffle(&mut v[..]);
        sort(&mut v[..]);
    });
    b.bytes = bytes as u64;
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! A handler and associated types for the filesystems action.
//!
//! The filesystems action lists all mounted filesystems on the client,
//! collecting device name, mount point, filesystem type and its options.
//! Current implementation works only in Linux systems.

use rrg_proto::{Filesystem, KeyValue, AttributedDict, DataBlob};
use crate::session::{self, Session};

use log::error;
use std::fmt::{Display, Formatter};

/// Enum of possible errors, which can occur during collecting filesystems data.
#[derive(Debug)]
enum Error {
    /// Missing mtab-like file error.
    MissingFile(std::io::Error),
    /// Parsing mtab-like file error.
    MountInfoParse(std::io::Error),
}

impl std::error::Error for Error {

    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        use Error::*;

        match *self {
            MissingFile(ref error) => Some(error),
            MountInfoParse(ref error) => Some(error),
        }
    }
}

impl Display for Error {

    fn fmt(&self, fmt: &mut Formatter) -> std::fmt::Result {
        use Error::*;

        match *self {
            MissingFile(ref error) => {
                write!(fmt, "missed file error: {}", error)
            },
            MountInfoParse(ref error) => {
                write!(fmt, "failed to obtain mount information: {}", error)
            },
        }
    }
}

impl From<Error> for session::Error {

    fn from(error: Error) -> session::Error {
        session::Error::action(error)
    }
}

/// A response type for the filesystems action.
pub struct Response {
    /// Information about the filesystem.
    mount_info: proc_mounts::MountInfo,
}

/// Handles requests for the filesystems action.
///
/// Initially searches in `/proc/mounts`. If it's missing, falls back to
/// `/etc/mtab`.
pub fn handle<S: Session>(session: &mut S, _: ()) -> session::Result<()> {
    use proc_mounts::MountIter;

    // Check whether `/proc/mounts` exists.
    let mount_iter = match MountIter::new() {
        Ok(mount_iter) => mount_iter,
        Err(_) => {
            // `/proc/mounts` doesn't exist. Try to fall back to `/etc/mtab`.
            MountIter::new_from_file("/etc/mtab").map_err(Error::MissingFile)?
        },
    };

    for mount_info in mount_iter {
        let mount_info = mount_info.map_err(Error::MountInfoParse)?;
        session.reply(Response {
           mount_info,
        })?;
    }

    Ok(())
}

/// Converts filesystem mount option in `String` representation to
/// GRR's `KeyValue` protobuf struct representation.
fn option_to_key_value(option: String) -> KeyValue {
    match &option.split('=').collect::<Vec<&str>>()[..] {
        &[key] => {
            // TODO: Simplify work with `DataBlob`.
            KeyValue {
                k: Some(DataBlob {
                    string: Some(String::from(key)),
                    ..Default::default()
                }),
                v: None,
            }
        },
        &[key, value] => {
            KeyValue {
                k: Some(DataBlob {
                    string: Some(String::from(key)),
                    ..Default::default()
                }),
                v: Some(DataBlob {
                    string: Some(String::from(value)),
                    ..Default::default()
                }),
            }
        },
        _ => {
            error!("invalid mount option syntax: {}", option);
            KeyValue {
                k: None,
                v: None,
            }
        },
    }
}

/// Converts a `Vec` of filesystem mount options in `String` representation to
/// GRR's `AttributedDict` protobuf struct representation.
fn options_to_dict(options: Vec<String>) -> AttributedDict {
    AttributedDict {
        dat: options.into_iter().map(option_to_key_value).collect(),
    }
}

impl super::Response for Response {

    const RDF_NAME: Option<&'static str> = Some("Filesystem");

    type Proto = rrg_proto::Filesystem;

    fn into_proto(self) -> Filesystem {
        // TODO: Remove lossy conversion of `PathBuf` to `String`
        // when `mount_point` and `device` fields of `Filesystem` message
        // will have `bytes` type instead of `string`.
        Filesystem {
            device: Some(self.mount_info.source.to_string_lossy()
                .into_owned()),
            mount_point: Some(self.mount_info.dest.to_string_lossy()
                .into_owned()),
            r#type: Some(self.mount_info.fstype),
            label: None,
            options: Some(options_to_dict(self.mount_info.options)),
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use std::path::PathBuf;

    #[test]
    fn test_if_any_filesystem_exists() {
        let mut session = session::test::Fake::new();
        assert!(handle(&mut session, ()).is_ok());
        assert_ne!(session.reply_count(), 0);
    }

    /// Returns all responses whose `source` is equal to `fs_name`.
    fn find_filesystems_by_name<'a>(
        session: &'a session::test::Fake,
        fs_name: &'a PathBuf,
    ) -> Vec<&'a Response> {
        let mut responses = Vec::new();

        for i in 0..session.reply_count() {
            let response: &Response = session.reply(i);
            if response.mount_info.source == *fs_name {
                responses.push(response);
            }
        }

        responses
    }

    /// Unit-like struct, representing a filesystem for testing with `fuse`.
    struct FuseFilesystem;

    impl fuse::Filesystem for FuseFilesystem {}

    #[test]
    fn test_fuse_filesystem() {
        use users::{get_current_uid, get_current_gid};

        let fs_name = PathBuf::from("fuse-test-fs");

        let tmp_dir = tempfile::tempdir().unwrap();
        let mountpoint = tmp_dir.path();

        let fs_name_option = format!("fsname={}", fs_name.to_str().unwrap());
        let options = [
            "-o", "ro",
            "-o", "nosuid",
            "-o", "nodev",
            "-o", "relatime",
            "-o", "subtype=custom-type",
            "-o", &fs_name_option,
        ];
        let options = options.iter().map(|opt| opt.as_ref())
            .collect::<Vec<&std::ffi::OsStr>>();

        // Spawn a background thread to handle filesystem operations.
        // When `_background_session` is dropped, filesystem will be unmounted.
        let background_session = unsafe {
            fuse::spawn_mount(FuseFilesystem, &mountpoint, &options).unwrap()
        };

        let mut session = session::test::Fake::new();
        assert!(handle(&mut session, ()).is_ok());

        let fuse_mounted_fs = find_filesystems_by_name(&session, &fs_name);
        assert_eq!(fuse_mounted_fs.len(), 1);

        let mount_info = &fuse_mounted_fs[0].mount_info;
        assert_eq!(mount_info.source, fs_name);
        assert_eq!(mount_info.dest, tmp_dir.path());
        assert_eq!(mount_info.fstype, "fuse.custom-type");

        let options = &mount_info.options;
        assert!(options.iter().any(|opt| opt == "ro"));
        assert!(options.iter().any(|opt| opt == "nosuid"));
        assert!(options.iter().any(|opt| opt == "nodev"));
        assert!(options.iter().any(|opt| opt == "relatime"));

        // `user_id` and `group_id` are set by libfuse.
        // http://man7.org/linux/man-pages/man8/mount.fuse.8.html
        let current_uid_option = format!("user_id={}", get_current_uid());
        assert!(options.iter().any(|opt| *opt == current_uid_option));

        let current_gid_option = format!("group_id={}", get_current_gid());
        assert!(options.iter().any(|opt| *opt == current_gid_option));

        drop(background_session);
    }
}
<commit_msg>Refactor tests for the filesystems action.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! A handler and associated types for the filesystems action.
//!
//! The filesystems action lists all mounted filesystems on the client,
//! collecting device name, mount point, filesystem type and its options.
//! Current implementation works only in Linux systems.

use rrg_proto::{Filesystem, KeyValue, AttributedDict, DataBlob};
use crate::session::{self, Session};

use log::error;
use std::fmt::{Display, Formatter};

/// Enum of possible errors, which can occur during collecting filesystems data.
#[derive(Debug)]
enum Error {
    /// Missing mtab-like file error.
    MissingFile(std::io::Error),
    /// Parsing mtab-like file error.
    MountInfoParse(std::io::Error),
}

impl std::error::Error for Error {

    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        use Error::*;

        match *self {
            MissingFile(ref error) => Some(error),
            MountInfoParse(ref error) => Some(error),
        }
    }
}

impl Display for Error {

    fn fmt(&self, fmt: &mut Formatter) -> std::fmt::Result {
        use Error::*;

        match *self {
            MissingFile(ref error) => {
                write!(fmt, "missed file error: {}", error)
            },
            MountInfoParse(ref error) => {
                write!(fmt, "failed to obtain mount information: {}", error)
            },
        }
    }
}

impl From<Error> for session::Error {

    fn from(error: Error) -> session::Error {
        session::Error::action(error)
    }
}

/// A response type for the filesystems action.
pub struct Response {
    /// Information about the filesystem.
    mount_info: proc_mounts::MountInfo,
}

/// Handles requests for the filesystems action.
///
/// Initially searches in `/proc/mounts`. If it's missing, falls back to
/// `/etc/mtab`.
pub fn handle<S: Session>(session: &mut S, _: ()) -> session::Result<()> {
    use proc_mounts::MountIter;

    // Check whether `/proc/mounts` exists.
    let mount_iter = match MountIter::new() {
        Ok(mount_iter) => mount_iter,
        Err(_) => {
            // `/proc/mounts` doesn't exist. Try to fall back to `/etc/mtab`.
            MountIter::new_from_file("/etc/mtab").map_err(Error::MissingFile)?
        },
    };

    for mount_info in mount_iter {
        let mount_info = mount_info.map_err(Error::MountInfoParse)?;
        session.reply(Response {
           mount_info,
        })?;
    }

    Ok(())
}

/// Converts filesystem mount option in `String` representation to
/// GRR's `KeyValue` protobuf struct representation.
fn option_to_key_value(option: String) -> KeyValue {
    match &option.split('=').collect::<Vec<&str>>()[..] {
        &[key] => {
            // TODO: Simplify work with `DataBlob`.
            KeyValue {
                k: Some(DataBlob {
                    string: Some(String::from(key)),
                    ..Default::default()
                }),
                v: None,
            }
        },
        &[key, value] => {
            KeyValue {
                k: Some(DataBlob {
                    string: Some(String::from(key)),
                    ..Default::default()
                }),
                v: Some(DataBlob {
                    string: Some(String::from(value)),
                    ..Default::default()
                }),
            }
        },
        _ => {
            error!("invalid mount option syntax: {}", option);
            KeyValue {
                k: None,
                v: None,
            }
        },
    }
}

/// Converts a `Vec` of filesystem mount options in `String` representation to
/// GRR's `AttributedDict` protobuf struct representation.
fn options_to_dict(options: Vec<String>) -> AttributedDict {
    AttributedDict {
        dat: options.into_iter().map(option_to_key_value).collect(),
    }
}

impl super::Response for Response {

    const RDF_NAME: Option<&'static str> = Some("Filesystem");

    type Proto = rrg_proto::Filesystem;

    fn into_proto(self) -> Filesystem {
        // TODO: Remove lossy conversion of `PathBuf` to `String`
        // when `mount_point` and `device` fields of `Filesystem` message
        // will have `bytes` type instead of `string`.
        Filesystem {
            device: Some(self.mount_info.source.to_string_lossy()
                .into_owned()),
            mount_point: Some(self.mount_info.dest.to_string_lossy()
                .into_owned()),
            r#type: Some(self.mount_info.fstype),
            label: None,
            options: Some(options_to_dict(self.mount_info.options)),
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use std::path::PathBuf;

    #[test]
    fn test_if_any_filesystem_exists() {
        let mut session = session::test::Fake::new();
        assert!(handle(&mut session, ()).is_ok());
        assert_ne!(session.reply_count(), 0);
    }

    /// Unit-like struct, representing a filesystem for testing with `fuse`.
    struct FuseFilesystem;

    impl fuse::Filesystem for FuseFilesystem {}

    #[test]
    fn test_fuse_filesystem() {
        use users::{get_current_uid, get_current_gid};

        let fs_name = PathBuf::from("fuse-test-fs");

        let tmp_dir = tempfile::tempdir().unwrap();
        let mountpoint = tmp_dir.path();

        let fs_name_option = format!("fsname={}", fs_name.to_str().unwrap());
        let options = [
            "-o", "ro",
            "-o", "nosuid",
            "-o", "nodev",
            "-o", "relatime",
            "-o", "subtype=custom-type",
            "-o", &fs_name_option,
        ];
        let options = options.iter().map(|opt| opt.as_ref())
            .collect::<Vec<&std::ffi::OsStr>>();

        // Spawn a background thread to handle filesystem operations.
        // When `_background_session` is dropped, filesystem will be unmounted.
        let background_session = unsafe {
            fuse::spawn_mount(FuseFilesystem, &mountpoint, &options).unwrap()
        };

        let mut session = session::test::Fake::new();
        assert!(handle(&mut session, ()).is_ok());

        let fuse_mounted_fs = session.replies::<Response>()
            .find(|reply| reply.mount_info.source == fs_name)
            .expect("no reply with a mounted FUSE filesystem");

        let mount_info = &fuse_mounted_fs.mount_info;
        assert_eq!(mount_info.source, fs_name);
        assert_eq!(mount_info.dest, tmp_dir.path());
        assert_eq!(mount_info.fstype, "fuse.custom-type");

        let options = &mount_info.options;
        assert!(options.iter().any(|opt| opt == "ro"));
        assert!(options.iter().any(|opt| opt == "nosuid"));
        assert!(options.iter().any(|opt| opt == "nodev"));
        assert!(options.iter().any(|opt| opt == "relatime"));

        // `user_id` and `group_id` are set by libfuse.
        // http://man7.org/linux/man-pages/man8/mount.fuse.8.html
        let current_uid_option = format!("user_id={}", get_current_uid());
        assert!(options.iter().any(|opt| *opt == current_uid_option));

        let current_gid_option = format!("group_id={}", get_current_gid());
        assert!(options.iter().any(|opt| *opt == current_gid_option));

        drop(background_session);
    }
}
<|endoftext|>


<commit_before>//! Virtual address space definitions

use util;
use ::core::marker::PhantomData;
use ::core::fmt::Debug;
use ::core::ops::Deref;
use vspace::VSpaceWindow;
use types::*;

/// The low boot window is a 1-1 mapped 4GB window of the bottom of memory
/// This window is used both as where the boot code initially runs before
/// setting up virtual addresses, and is kept through some of the early
/// kernel init to access boot data structures.
/// We define this window starting at 1 and being 1 byte less in size to
/// prevent the creation of objects at address 0, which would be undefined
/// behaviour
const LOW_BOOT_MAPPING: (usize, usize) = (0x0 + 1, util::GB * 4 - 1);

/// Early high window is a slice of the final kernel window. It covers a
/// single gigabyte, which is meant to be enough to allow the kernel code
/// to be executed at it's final address. This region is also mapped
/// to the same (first 1gb) of the low boot window
const HIGH_BOOT_MAPPING: (usize, usize) = (0xffffffff80000000, util::GB);

/// Final kernel window is the top 2^39 bits of memory
const KERNEL_MAPPING: (usize, usize) = (0xffffff8000000000, 0x8000000000);

/// The low window should should only be constructed immediately on boot
/// entry, and then dropped before switching away from the bootstrapping
/// address space
pub struct BootLowWindow<'a>(PhantomData<&'a usize>);
/// The high kernel window is valid initially on boot, and should always stay
/// valid. As a result the high window should be constructed as a `subwindow`
/// of the `KernelWindow`
pub struct BootHighWindow<'a>(PhantomData<&'a usize>);
/// The final kernel window is valid after early boot has happened, and
/// remains valid forever after
pub struct KernelWindow<'a>(PhantomData<&'a usize>);

/// Wrapper for an address in a high window
#[derive(Ord, Eq, PartialEq, PartialOrd, Debug, Copy, Clone)]
pub struct HighWindowAddr(usize);

/// Wrapper for an address in the kernel window
#[derive(Ord, Eq, PartialEq, PartialOrd, Debug, Copy, Clone)]
pub struct KernelWindowAddr(usize);

/// Wrapper for an address in the boot window
#[derive(Ord, Eq, PartialEq, PartialOrd, Debug, Copy, Clone)]
pub struct LowWindowAddr(usize);

impl Deref for HighWindowAddr{
    type Target = usize;
    fn deref(&self) -> &usize {
        &self.0
    }
}

impl Deref for LowWindowAddr{
    type Target = usize;
    fn deref(&self) -> &usize {
        &self.0
    }
}

impl Deref for KernelWindowAddr{
    type Target = usize;
    fn deref(&self) -> &usize {
        &self.0
    }
}

unsafe impl<'a> VSpaceWindow<'a> for BootHighWindow<'a> {
    type Addr = HighWindowAddr;
    fn base(&self) -> usize { HIGH_BOOT_MAPPING.0 }
    fn size(&self) -> usize { HIGH_BOOT_MAPPING.1 }
    unsafe fn default() -> BootHighWindow<'a> {
        BootHighWindow(PhantomData)
    }
    unsafe fn to_paddr(&self, addr: Self::Addr) -> PAddr {
        PAddr(addr.0 - HIGH_BOOT_MAPPING.0)
    }
    unsafe fn from_paddr(&self, paddr: PAddr) -> Self::Addr {
        HighWindowAddr(paddr.0 + HIGH_BOOT_MAPPING.0)
    }
    unsafe fn to_addr(&self, addr: usize) -> Self::Addr {
        HighWindowAddr(addr)
    }
}

unsafe impl<'a> VSpaceWindow<'a> for KernelWindow<'a> {
    type Addr = KernelWindowAddr;
    fn base(&self) -> usize { KERNEL_MAPPING.0 }
    fn size(&self) -> usize { KERNEL_MAPPING.1 }
    unsafe fn default() -> KernelWindow<'a> {
        KernelWindow(PhantomData)
    }
    unsafe fn to_paddr(&self, addr: Self::Addr) -> PAddr {
        unimplemented!()
    }
    unsafe fn from_paddr(&self, paddr: PAddr) -> Self::Addr {
        unimplemented!()
    }
    unsafe fn to_addr(&self, addr: usize) -> Self::Addr {
        KernelWindowAddr(addr)
    }
}

unsafe impl<'a> VSpaceWindow<'a> for BootLowWindow<'a> {
    type Addr = LowWindowAddr;
    fn base(&self) -> usize { LOW_BOOT_MAPPING.0 }
    fn size(&self) -> usize { LOW_BOOT_MAPPING.1 }
    unsafe fn default() -> BootLowWindow<'a> {
        BootLowWindow(PhantomData)
    }
    unsafe fn to_paddr(&self, addr: Self::Addr) -> PAddr {
        PAddr(addr.0)
    }
    unsafe fn from_paddr(&self, paddr: PAddr) -> Self::Addr {
        LowWindowAddr(paddr.0)
    }
    unsafe fn to_addr(&self, addr: usize) -> Self::Addr {
        LowWindowAddr(addr)
    }
}
<commit_msg>Add some debug asserts on vspace windows<commit_after>//! Virtual address space definitions

use util;
use ::core::marker::PhantomData;
use ::core::fmt::Debug;
use ::core::ops::Deref;
use vspace::VSpaceWindow;
use types::*;

/// The low boot window is a 1-1 mapped 4GB window of the bottom of memory
/// This window is used both as where the boot code initially runs before
/// setting up virtual addresses, and is kept through some of the early
/// kernel init to access boot data structures.
/// We define this window starting at 1 and being 1 byte less in size to
/// prevent the creation of objects at address 0, which would be undefined
/// behaviour
const LOW_BOOT_MAPPING: (usize, usize) = (0x0 + 1, util::GB * 4 - 1);

/// Early high window is a slice of the final kernel window. It covers a
/// single gigabyte, which is meant to be enough to allow the kernel code
/// to be executed at it's final address. This region is also mapped
/// to the same (first 1gb) of the low boot window
const HIGH_BOOT_MAPPING: (usize, usize) = (0xffffffff80000000, util::GB);

/// Final kernel window is the top 2^39 bits of memory
const KERNEL_MAPPING: (usize, usize) = (0xffffff8000000000, 0x8000000000);

/// The low window should should only be constructed immediately on boot
/// entry, and then dropped before switching away from the bootstrapping
/// address space
pub struct BootLowWindow<'a>(PhantomData<&'a usize>);
/// The high kernel window is valid initially on boot, and should always stay
/// valid. As a result the high window should be constructed as a `subwindow`
/// of the `KernelWindow`
pub struct BootHighWindow<'a>(PhantomData<&'a usize>);
/// The final kernel window is valid after early boot has happened, and
/// remains valid forever after
pub struct KernelWindow<'a>(PhantomData<&'a usize>);

/// Wrapper for an address in a high window
#[derive(Ord, Eq, PartialEq, PartialOrd, Debug, Copy, Clone)]
pub struct HighWindowAddr(usize);

/// Wrapper for an address in the kernel window
#[derive(Ord, Eq, PartialEq, PartialOrd, Debug, Copy, Clone)]
pub struct KernelWindowAddr(usize);

/// Wrapper for an address in the boot window
#[derive(Ord, Eq, PartialEq, PartialOrd, Debug, Copy, Clone)]
pub struct LowWindowAddr(usize);

impl Deref for HighWindowAddr{
    type Target = usize;
    fn deref(&self) -> &usize {
        &self.0
    }
}

impl Deref for LowWindowAddr{
    type Target = usize;
    fn deref(&self) -> &usize {
        &self.0
    }
}

impl Deref for KernelWindowAddr{
    type Target = usize;
    fn deref(&self) -> &usize {
        &self.0
    }
}

unsafe impl<'a> VSpaceWindow<'a> for BootHighWindow<'a> {
    type Addr = HighWindowAddr;
    fn base(&self) -> usize { HIGH_BOOT_MAPPING.0 }
    fn size(&self) -> usize { HIGH_BOOT_MAPPING.1 }
    unsafe fn default() -> BootHighWindow<'a> {
        BootHighWindow(PhantomData)
    }
    unsafe fn to_paddr(&self, addr: Self::Addr) -> PAddr {
        debug_assert!(self.addr_range_valid(addr, 0));
        PAddr(addr.0 - HIGH_BOOT_MAPPING.0)
    }
    unsafe fn from_paddr(&self, paddr: PAddr) -> Self::Addr {
        self.to_addr(paddr.0 + HIGH_BOOT_MAPPING.0)
    }
    unsafe fn to_addr(&self, addr: usize) -> Self::Addr {
        debug_assert!(self.range_valid(addr, 0));
        HighWindowAddr(addr)
    }
}

unsafe impl<'a> VSpaceWindow<'a> for KernelWindow<'a> {
    type Addr = KernelWindowAddr;
    fn base(&self) -> usize { KERNEL_MAPPING.0 }
    fn size(&self) -> usize { KERNEL_MAPPING.1 }
    unsafe fn default() -> KernelWindow<'a> {
        KernelWindow(PhantomData)
    }
    unsafe fn to_paddr(&self, addr: Self::Addr) -> PAddr {
        unimplemented!()
    }
    unsafe fn from_paddr(&self, paddr: PAddr) -> Self::Addr {
        unimplemented!()
    }
    unsafe fn to_addr(&self, addr: usize) -> Self::Addr {
        unimplemented!()
    }
}

unsafe impl<'a> VSpaceWindow<'a> for BootLowWindow<'a> {
    type Addr = LowWindowAddr;
    fn base(&self) -> usize { LOW_BOOT_MAPPING.0 }
    fn size(&self) -> usize { LOW_BOOT_MAPPING.1 }
    unsafe fn default() -> BootLowWindow<'a> {
        BootLowWindow(PhantomData)
    }
    unsafe fn to_paddr(&self, addr: Self::Addr) -> PAddr {
        debug_assert!(self.addr_range_valid(addr, 0));
        PAddr(addr.0)
    }
    unsafe fn from_paddr(&self, paddr: PAddr) -> Self::Addr {
        self.to_addr(paddr.0)
    }
    unsafe fn to_addr(&self, addr: usize) -> Self::Addr {
        debug_assert!(self.range_valid(addr, 0));
        LowWindowAddr(addr)
    }
}
<|endoftext|>