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

//! Protocol used to ensure peer liveness
//!
//! The HealthChecker is responsible for ensuring liveness of all peers of a node.
//! It does so by periodically selecting a random connected peer and sending a Ping probe. A
//! healthy peer is expected to respond with a corresponding Pong message.
//!
//! If a certain number of successive liveness probes for a peer fail, the HealthChecker initiates a
//! disconnect from the peer. It relies on ConnectivityManager or the remote peer to re-establish
//! the connection.
//!
//! Future Work
//! -----------
//! We can make a few other improvements to the health checker. These are:
//! - Make the policy for interpreting ping failures pluggable
//! - Use successful inbound pings as a sign of remote note being healthy
//! - Ping a peer only in periods of no application-level communication with the peer
use crate::{
    application::interface::NetworkInterface,
    constants::NETWORK_CHANNEL_SIZE,
    counters,
    error::NetworkError,
    logging::NetworkSchema,
    peer_manager::{ConnectionRequestSender, PeerManagerRequestSender},
    protocols::{
        health_checker::interface::{HealthCheckData, HealthCheckNetworkInterface},
        network::{
            AppConfig, ApplicationNetworkSender, Event, NetworkEvents, NetworkSender,
            NewNetworkSender,
        },
        rpc::error::RpcError,
    },
    ProtocolId,
};
use aptos_config::network_id::{NetworkContext, PeerNetworkId};
use aptos_logger::prelude::*;
use aptos_time_service::{TimeService, TimeServiceTrait};
use aptos_types::PeerId;
use async_trait::async_trait;
use bytes::Bytes;
use channel::{aptos_channel, message_queues::QueueStyle};
use futures::{
    channel::oneshot,
    stream::{FuturesUnordered, StreamExt},
};
use rand::{rngs::SmallRng, Rng, SeedableRng};
use serde::{Deserialize, Serialize};
use short_hex_str::AsShortHexStr;
use std::{collections::hash_map::Entry, time::Duration};

pub mod builder;
mod interface;
#[cfg(test)]
mod test;

/// The interface from Network to HealthChecker layer.
///
/// `HealthCheckerNetworkEvents` is a `Stream` of `PeerManagerNotification` where the
/// raw `Bytes` rpc messages are deserialized into
/// `HealthCheckerMsg` types. `HealthCheckerNetworkEvents` is a thin wrapper
/// around an `channel::Receiver<PeerManagerNotification>`.
pub type HealthCheckerNetworkEvents = NetworkEvents<HealthCheckerMsg>;

/// The interface from HealthChecker to Networking layer.
///
/// This is a thin wrapper around a `NetworkSender<HealthCheckerMsg>`, so it is
/// easy to clone and send off to a separate task. For example, the rpc requests
/// return Futures that encapsulate the whole flow, from sending the request to
/// remote, to finally receiving the response and deserializing. It therefore
/// makes the most sense to make the rpc call on a separate async task, which
/// requires the `HealthCheckerNetworkSender` to be `Clone` and `Send`.
#[derive(Clone)]
pub struct HealthCheckerNetworkSender {
    inner: NetworkSender<HealthCheckerMsg>,
}

/// Configuration for the network endpoints to support HealthChecker.
pub fn network_endpoint_config() -> AppConfig {
    AppConfig::p2p(
        [ProtocolId::HealthCheckerRpc],
        aptos_channel::Config::new(NETWORK_CHANNEL_SIZE)
            .queue_style(QueueStyle::LIFO)
            .counters(&counters::PENDING_HEALTH_CHECKER_NETWORK_EVENTS),
    )
}

impl NewNetworkSender for HealthCheckerNetworkSender {
    fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self {
        Self {
            inner: NetworkSender::new(peer_mgr_reqs_tx, connection_reqs_tx),
        }
    }
}

impl HealthCheckerNetworkSender {
    pub async fn disconnect_peer(&mut self, peer_id: PeerId) -> Result<(), NetworkError> {
        self.inner.disconnect_peer(peer_id).await
    }
}

#[async_trait]
impl ApplicationNetworkSender<HealthCheckerMsg> for HealthCheckerNetworkSender {
    /// Send a HealthChecker Ping RPC request to remote peer `recipient`. Returns
    /// the remote peer's future `Pong` reply.
    ///
    /// The rpc request can be canceled at any point by dropping the returned
    /// future.
    async fn send_rpc(
        &self,
        recipient: PeerId,
        req_msg: HealthCheckerMsg,
        timeout: Duration,
    ) -> Result<HealthCheckerMsg, RpcError> {
        let protocol = ProtocolId::HealthCheckerRpc;
        self.inner
            .send_rpc(recipient, protocol, req_msg, timeout)
            .await
    }
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub enum HealthCheckerMsg {
    Ping(Ping),
    Pong(Pong),
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Ping(u32);

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Pong(u32);

/// The actor performing health checks by running the Ping protocol
pub struct HealthChecker {
    network_context: NetworkContext,
    /// A handle to a time service for easily mocking time-related operations.
    time_service: TimeService,
    /// Network interface to send requests to the Network Layer
    network_interface: HealthCheckNetworkInterface,
    /// Random-number generator.
    rng: SmallRng,
    /// Time we wait between each set of pings.
    ping_interval: Duration,
    /// Ping timeout duration.
    ping_timeout: Duration,
    /// Number of successive ping failures we tolerate before declaring a node as unhealthy and
    /// disconnecting from it. In the future, this can be replaced with a more general failure
    /// detection policy.
    ping_failures_tolerated: u64,
    /// Counter incremented in each round of health checks
    round: u64,
}

impl HealthChecker {
    /// Create new instance of the [`HealthChecker`] actor.
    pub fn new(
        network_context: NetworkContext,
        time_service: TimeService,
        network_interface: HealthCheckNetworkInterface,
        ping_interval: Duration,
        ping_timeout: Duration,
        ping_failures_tolerated: u64,
    ) -> Self {
        HealthChecker {
            network_context,
            time_service,
            network_interface,
            rng: SmallRng::from_entropy(),
            ping_interval,
            ping_timeout,
            ping_failures_tolerated,
            round: 0,
        }
    }

    pub async fn start(mut self) {
        let mut tick_handlers = FuturesUnordered::new();
        info!(
            NetworkSchema::new(&self.network_context),
            "{} Health checker actor started", self.network_context
        );

        let ticker = self.time_service.interval(self.ping_interval);
        tokio::pin!(ticker);

        loop {
            futures::select! {
                maybe_event = self.network_interface.next() => {
                    // Shutdown the HealthChecker when this network instance shuts
                    // down. This happens when the `PeerManager` drops.
                    let event = match maybe_event {
                        Some(event) => event,
                        None => break,
                    };

                    match event {
                        Event::NewPeer(metadata) => {
                            self.network_interface.app_data().insert(
                                metadata.remote_peer_id,
                                HealthCheckData::new(self.round)
                            );
                        }
                        Event::LostPeer(metadata) => {
                            self.network_interface.app_data().remove(
                                &metadata.remote_peer_id
                            );
                        }
                        Event::RpcRequest(peer_id, msg, protocol, res_tx) => {
                            match msg {
                                HealthCheckerMsg::Ping(ping) => self.handle_ping_request(peer_id, ping, protocol, res_tx),
                                _ => {
                                    warn!(
                                        SecurityEvent::InvalidHealthCheckerMsg,
                                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                                        rpc_message = msg,
                                        "{} Unexpected RPC message from {}",
                                        self.network_context,
                                        peer_id
                                    );
                                    debug_assert!(false, "Unexpected rpc request");
                                }
                            };
                        }
                        Event::Message(peer_id, msg) => {
                            error!(
                                SecurityEvent::InvalidNetworkEventHC,
                                NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                                "{} Unexpected direct send from {} msg {:?}",
                                self.network_context,
                                peer_id,
                                msg,
                            );
                            debug_assert!(false, "Unexpected network event");
                        }
                    }
                }
                _ = ticker.select_next_some() => {
                    self.round += 1;
                    let connected = self.network_interface.connected_peers();
                    if connected.is_empty() {
                        trace!(
                            NetworkSchema::new(&self.network_context),
                            round = self.round,
                            "{} No connected peer to ping round: {}",
                            self.network_context,
                            self.round
                        );
                        continue
                    }

                    for peer_id in connected {
                        let nonce = self.rng.gen::<u32>();
                        trace!(
                            NetworkSchema::new(&self.network_context),
                            round = self.round,
                            "{} Will ping: {} for round: {} nonce: {}",
                            self.network_context,
                            peer_id.short_str(),
                            self.round,
                            nonce
                        );

                        tick_handlers.push(Self::ping_peer(
                            self.network_context,
                            self.network_interface.sender(),
                            peer_id,
                            self.round,
                            nonce,
                            self.ping_timeout,
                        ));
                    }
                }
                res = tick_handlers.select_next_some() => {
                    let (peer_id, round, nonce, ping_result) = res;
                    self.handle_ping_response(peer_id, round, nonce, ping_result).await;
                }
            }
        }
        warn!(
            NetworkSchema::new(&self.network_context),
            "{} Health checker actor terminated", self.network_context
        );
    }

    fn handle_ping_request(
        &mut self,
        peer_id: PeerId,
        ping: Ping,
        protocol: ProtocolId,
        res_tx: oneshot::Sender<Result<Bytes, RpcError>>,
    ) {
        let message = match protocol.to_bytes(&HealthCheckerMsg::Pong(Pong(ping.0))) {
            Ok(msg) => msg,
            Err(e) => {
                warn!(
                    NetworkSchema::new(&self.network_context),
                    error = ?e,
                    "{} Unable to serialize pong response: {}", self.network_context, e
                );
                return;
            }
        };
        trace!(
            NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
            "{} Sending Pong response to peer: {} with nonce: {}",
            self.network_context,
            peer_id.short_str(),
            ping.0,
        );
        let _ = res_tx.send(Ok(message.into()));
    }

    async fn handle_ping_response(
        &mut self,
        peer_id: PeerId,
        round: u64,
        req_nonce: u32,
        ping_result: Result<Pong, RpcError>,
    ) {
        match ping_result {
            Ok(pong) => {
                if pong.0 == req_nonce {
                    trace!(
                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                        rount = round,
                        "{} Ping successful for peer: {} round: {}",
                        self.network_context,
                        peer_id.short_str(),
                        round
                    );
                    // Update last successful ping to current round.
                    // If it's not in storage, don't bother updating it
                    let _ = self.network_interface.app_data().write(peer_id, |entry| {
                        match entry {
                            Entry::Vacant(..) => {
                                // Don't do anything if there isn't an entry
                            }
                            Entry::Occupied(inner) => {
                                let data = inner.get_mut();
                                // Update state if it's a newer round
                                if round > data.round {
                                    data.round = round;
                                    data.failures = 0;
                                }
                            }
                        };
                        Ok(())
                    });
                } else {
                    warn!(
                        SecurityEvent::InvalidHealthCheckerMsg,
                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                        "{} Pong nonce doesn't match Ping nonce. Round: {}, Pong: {}, Ping: {}",
                        self.network_context,
                        round,
                        pong.0,
                        req_nonce
                    );
                    debug_assert!(false, "Pong nonce doesn't match our challenge Ping nonce");
                }
            }
            Err(err) => {
                warn!(
                    NetworkSchema::new(&self.network_context)
                        .remote_peer(&peer_id),
                    error = ?err,
                    round = round,
                    "{} Ping failed for peer: {} round: {} with error: {:?}",
                    self.network_context,
                    peer_id.short_str(),
                    round,
                    err
                );
                let _ = self.network_interface.app_data().write(peer_id, |entry| {
                    // Don't add in a default in case it's already disconnected
                    match entry {
                        Entry::Vacant(..) => {
                            // Don't do anything if there isn't an entry
                        }
                        Entry::Occupied(inner) => {
                            // If this is the result of an older ping, we ignore it.
                            // Increment num of failures.
                            let data = inner.get_mut();
                            if data.round <= round {
                                data.failures += 1;
                            }
                        }
                    }
                    Ok(())
                });

                // If the ping failures are now more than
                // `self.ping_failures_tolerated`, we disconnect from the node.
                // The HealthChecker only performs the disconnect. It relies on
                // ConnectivityManager or the remote peer to re-establish the connection.
                let failures = self
                    .network_interface
                    .app_data()
                    .read(&peer_id)
                    .map(|data| data.failures)
                    .unwrap_or(0);
                if failures > self.ping_failures_tolerated {
                    info!(
                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                        "{} Disconnecting from peer: {}",
                        self.network_context,
                        peer_id.short_str()
                    );
                    let peer_network_id =
                        PeerNetworkId::new(self.network_context.network_id(), peer_id);
                    if let Err(err) = self
                        .network_interface
                        .disconnect_peer(peer_network_id)
                        .await
                    {
                        warn!(
                            NetworkSchema::new(&self.network_context)
                                .remote_peer(&peer_id),
                            error = ?err,
                            "{} Failed to disconnect from peer: {} with error: {:?}",
                            self.network_context,
                            peer_id.short_str(),
                            err
                        );
                    }
                }
            }
        }
    }

    async fn ping_peer(
        network_context: NetworkContext,
        network_tx: HealthCheckerNetworkSender,
        peer_id: PeerId,
        round: u64,
        nonce: u32,
        ping_timeout: Duration,
    ) -> (PeerId, u64, u32, Result<Pong, RpcError>) {
        trace!(
            NetworkSchema::new(&network_context).remote_peer(&peer_id),
            round = round,
            "{} Sending Ping request to peer: {} for round: {} nonce: {}",
            network_context,
            peer_id.short_str(),
            round,
            nonce
        );
        let res_pong_msg = network_tx
            .send_rpc(peer_id, HealthCheckerMsg::Ping(Ping(nonce)), ping_timeout)
            .await
            .and_then(|msg| match msg {
                HealthCheckerMsg::Pong(res) => Ok(res),
                _ => Err(RpcError::InvalidRpcResponse),
            });
        (peer_id, round, nonce, res_pong_msg)
    }
}
<commit_msg>[Networking HC] Reseting HC failure count on ingress HC request<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! Protocol used to ensure peer liveness
//!
//! The HealthChecker is responsible for ensuring liveness of all peers of a node.
//! It does so by periodically selecting a random connected peer and sending a Ping probe. A
//! healthy peer is expected to respond with a corresponding Pong message.
//!
//! If a certain number of successive liveness probes for a peer fail, the HealthChecker initiates a
//! disconnect from the peer. It relies on ConnectivityManager or the remote peer to re-establish
//! the connection.
//!
//! Future Work
//! -----------
//! We can make a few other improvements to the health checker. These are:
//! - Make the policy for interpreting ping failures pluggable
//! - Use successful inbound pings as a sign of remote note being healthy
//! - Ping a peer only in periods of no application-level communication with the peer
use crate::{
    application::interface::NetworkInterface,
    constants::NETWORK_CHANNEL_SIZE,
    counters,
    error::NetworkError,
    logging::NetworkSchema,
    peer_manager::{ConnectionRequestSender, PeerManagerRequestSender},
    protocols::{
        health_checker::interface::{HealthCheckData, HealthCheckNetworkInterface},
        network::{
            AppConfig, ApplicationNetworkSender, Event, NetworkEvents, NetworkSender,
            NewNetworkSender,
        },
        rpc::error::RpcError,
    },
    ProtocolId,
};
use aptos_config::network_id::{NetworkContext, PeerNetworkId};
use aptos_logger::prelude::*;
use aptos_time_service::{TimeService, TimeServiceTrait};
use aptos_types::PeerId;
use async_trait::async_trait;
use bytes::Bytes;
use channel::{aptos_channel, message_queues::QueueStyle};
use futures::{
    channel::oneshot,
    stream::{FuturesUnordered, StreamExt},
};
use rand::{rngs::SmallRng, Rng, SeedableRng};
use serde::{Deserialize, Serialize};
use short_hex_str::AsShortHexStr;
use std::{collections::hash_map::Entry, time::Duration};

pub mod builder;
mod interface;
#[cfg(test)]
mod test;

/// The interface from Network to HealthChecker layer.
///
/// `HealthCheckerNetworkEvents` is a `Stream` of `PeerManagerNotification` where the
/// raw `Bytes` rpc messages are deserialized into
/// `HealthCheckerMsg` types. `HealthCheckerNetworkEvents` is a thin wrapper
/// around an `channel::Receiver<PeerManagerNotification>`.
pub type HealthCheckerNetworkEvents = NetworkEvents<HealthCheckerMsg>;

/// The interface from HealthChecker to Networking layer.
///
/// This is a thin wrapper around a `NetworkSender<HealthCheckerMsg>`, so it is
/// easy to clone and send off to a separate task. For example, the rpc requests
/// return Futures that encapsulate the whole flow, from sending the request to
/// remote, to finally receiving the response and deserializing. It therefore
/// makes the most sense to make the rpc call on a separate async task, which
/// requires the `HealthCheckerNetworkSender` to be `Clone` and `Send`.
#[derive(Clone)]
pub struct HealthCheckerNetworkSender {
    inner: NetworkSender<HealthCheckerMsg>,
}

/// Configuration for the network endpoints to support HealthChecker.
pub fn network_endpoint_config() -> AppConfig {
    AppConfig::p2p(
        [ProtocolId::HealthCheckerRpc],
        aptos_channel::Config::new(NETWORK_CHANNEL_SIZE)
            .queue_style(QueueStyle::LIFO)
            .counters(&counters::PENDING_HEALTH_CHECKER_NETWORK_EVENTS),
    )
}

impl NewNetworkSender for HealthCheckerNetworkSender {
    fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self {
        Self {
            inner: NetworkSender::new(peer_mgr_reqs_tx, connection_reqs_tx),
        }
    }
}

impl HealthCheckerNetworkSender {
    pub async fn disconnect_peer(&mut self, peer_id: PeerId) -> Result<(), NetworkError> {
        self.inner.disconnect_peer(peer_id).await
    }
}

#[async_trait]
impl ApplicationNetworkSender<HealthCheckerMsg> for HealthCheckerNetworkSender {
    /// Send a HealthChecker Ping RPC request to remote peer `recipient`. Returns
    /// the remote peer's future `Pong` reply.
    ///
    /// The rpc request can be canceled at any point by dropping the returned
    /// future.
    async fn send_rpc(
        &self,
        recipient: PeerId,
        req_msg: HealthCheckerMsg,
        timeout: Duration,
    ) -> Result<HealthCheckerMsg, RpcError> {
        let protocol = ProtocolId::HealthCheckerRpc;
        self.inner
            .send_rpc(recipient, protocol, req_msg, timeout)
            .await
    }
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub enum HealthCheckerMsg {
    Ping(Ping),
    Pong(Pong),
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Ping(u32);

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Pong(u32);

/// The actor performing health checks by running the Ping protocol
pub struct HealthChecker {
    network_context: NetworkContext,
    /// A handle to a time service for easily mocking time-related operations.
    time_service: TimeService,
    /// Network interface to send requests to the Network Layer
    network_interface: HealthCheckNetworkInterface,
    /// Random-number generator.
    rng: SmallRng,
    /// Time we wait between each set of pings.
    ping_interval: Duration,
    /// Ping timeout duration.
    ping_timeout: Duration,
    /// Number of successive ping failures we tolerate before declaring a node as unhealthy and
    /// disconnecting from it. In the future, this can be replaced with a more general failure
    /// detection policy.
    ping_failures_tolerated: u64,
    /// Counter incremented in each round of health checks
    round: u64,
}

impl HealthChecker {
    /// Create new instance of the [`HealthChecker`] actor.
    pub fn new(
        network_context: NetworkContext,
        time_service: TimeService,
        network_interface: HealthCheckNetworkInterface,
        ping_interval: Duration,
        ping_timeout: Duration,
        ping_failures_tolerated: u64,
    ) -> Self {
        HealthChecker {
            network_context,
            time_service,
            network_interface,
            rng: SmallRng::from_entropy(),
            ping_interval,
            ping_timeout,
            ping_failures_tolerated,
            round: 0,
        }
    }

    pub async fn start(mut self) {
        let mut tick_handlers = FuturesUnordered::new();
        info!(
            NetworkSchema::new(&self.network_context),
            "{} Health checker actor started", self.network_context
        );

        let ticker = self.time_service.interval(self.ping_interval);
        tokio::pin!(ticker);

        loop {
            futures::select! {
                maybe_event = self.network_interface.next() => {
                    // Shutdown the HealthChecker when this network instance shuts
                    // down. This happens when the `PeerManager` drops.
                    let event = match maybe_event {
                        Some(event) => event,
                        None => break,
                    };

                    match event {
                        Event::NewPeer(metadata) => {
                            self.network_interface.app_data().insert(
                                metadata.remote_peer_id,
                                HealthCheckData::new(self.round)
                            );
                        }
                        Event::LostPeer(metadata) => {
                            self.network_interface.app_data().remove(
                                &metadata.remote_peer_id
                            );
                        }
                        Event::RpcRequest(peer_id, msg, protocol, res_tx) => {
                            match msg {
                                HealthCheckerMsg::Ping(ping) => self.handle_ping_request(peer_id, ping, protocol, res_tx),
                                _ => {
                                    warn!(
                                        SecurityEvent::InvalidHealthCheckerMsg,
                                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                                        rpc_message = msg,
                                        "{} Unexpected RPC message from {}",
                                        self.network_context,
                                        peer_id
                                    );
                                    debug_assert!(false, "Unexpected rpc request");
                                }
                            };
                        }
                        Event::Message(peer_id, msg) => {
                            error!(
                                SecurityEvent::InvalidNetworkEventHC,
                                NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                                "{} Unexpected direct send from {} msg {:?}",
                                self.network_context,
                                peer_id,
                                msg,
                            );
                            debug_assert!(false, "Unexpected network event");
                        }
                    }
                }
                _ = ticker.select_next_some() => {
                    self.round += 1;
                    let connected = self.network_interface.connected_peers();
                    if connected.is_empty() {
                        trace!(
                            NetworkSchema::new(&self.network_context),
                            round = self.round,
                            "{} No connected peer to ping round: {}",
                            self.network_context,
                            self.round
                        );
                        continue
                    }

                    for peer_id in connected {
                        let nonce = self.rng.gen::<u32>();
                        trace!(
                            NetworkSchema::new(&self.network_context),
                            round = self.round,
                            "{} Will ping: {} for round: {} nonce: {}",
                            self.network_context,
                            peer_id.short_str(),
                            self.round,
                            nonce
                        );

                        tick_handlers.push(Self::ping_peer(
                            self.network_context,
                            self.network_interface.sender(),
                            peer_id,
                            self.round,
                            nonce,
                            self.ping_timeout,
                        ));
                    }
                }
                res = tick_handlers.select_next_some() => {
                    let (peer_id, round, nonce, ping_result) = res;
                    self.handle_ping_response(peer_id, round, nonce, ping_result).await;
                }
            }
        }
        warn!(
            NetworkSchema::new(&self.network_context),
            "{} Health checker actor terminated", self.network_context
        );
    }

    fn handle_ping_request(
        &mut self,
        peer_id: PeerId,
        ping: Ping,
        protocol: ProtocolId,
        res_tx: oneshot::Sender<Result<Bytes, RpcError>>,
    ) {
        let message = match protocol.to_bytes(&HealthCheckerMsg::Pong(Pong(ping.0))) {
            Ok(msg) => msg,
            Err(e) => {
                warn!(
                    NetworkSchema::new(&self.network_context),
                    error = ?e,
                    "{} Unable to serialize pong response: {}", self.network_context, e
                );
                return;
            }
        };
        trace!(
            NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
            "{} Sending Pong response to peer: {} with nonce: {}",
            self.network_context,
            peer_id.short_str(),
            ping.0,
        );
        // Record Ingress HC here and reset failures.
        let _ = self.network_interface.app_data().write(peer_id, |entry| {
            match entry {
                Entry::Vacant(..) => {
                    // Don't do anything if there isn't an entry
                }
                Entry::Occupied(inner) => {
                    let data = inner.get_mut();
                    data.failures = 0;
                }
            };
            Ok(())
        });

        let _ = res_tx.send(Ok(message.into()));
    }

    async fn handle_ping_response(
        &mut self,
        peer_id: PeerId,
        round: u64,
        req_nonce: u32,
        ping_result: Result<Pong, RpcError>,
    ) {
        match ping_result {
            Ok(pong) => {
                if pong.0 == req_nonce {
                    trace!(
                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                        rount = round,
                        "{} Ping successful for peer: {} round: {}",
                        self.network_context,
                        peer_id.short_str(),
                        round
                    );
                    // Update last successful ping to current round.
                    // If it's not in storage, don't bother updating it
                    let _ = self.network_interface.app_data().write(peer_id, |entry| {
                        match entry {
                            Entry::Vacant(..) => {
                                // Don't do anything if there isn't an entry
                            }
                            Entry::Occupied(inner) => {
                                let data = inner.get_mut();
                                // Update state if it's a newer round
                                if round > data.round {
                                    data.round = round;
                                    data.failures = 0;
                                }
                            }
                        };
                        Ok(())
                    });
                } else {
                    warn!(
                        SecurityEvent::InvalidHealthCheckerMsg,
                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                        "{} Pong nonce doesn't match Ping nonce. Round: {}, Pong: {}, Ping: {}",
                        self.network_context,
                        round,
                        pong.0,
                        req_nonce
                    );
                    debug_assert!(false, "Pong nonce doesn't match our challenge Ping nonce");
                }
            }
            Err(err) => {
                warn!(
                    NetworkSchema::new(&self.network_context)
                        .remote_peer(&peer_id),
                    error = ?err,
                    round = round,
                    "{} Ping failed for peer: {} round: {} with error: {:?}",
                    self.network_context,
                    peer_id.short_str(),
                    round,
                    err
                );
                let _ = self.network_interface.app_data().write(peer_id, |entry| {
                    // Don't add in a default in case it's already disconnected
                    match entry {
                        Entry::Vacant(..) => {
                            // Don't do anything if there isn't an entry
                        }
                        Entry::Occupied(inner) => {
                            // If this is the result of an older ping, we ignore it.
                            // Increment num of failures.
                            let data = inner.get_mut();
                            if data.round <= round {
                                data.failures += 1;
                            }
                        }
                    }
                    Ok(())
                });

                // If the ping failures are now more than
                // `self.ping_failures_tolerated`, we disconnect from the node.
                // The HealthChecker only performs the disconnect. It relies on
                // ConnectivityManager or the remote peer to re-establish the connection.
                let failures = self
                    .network_interface
                    .app_data()
                    .read(&peer_id)
                    .map(|data| data.failures)
                    .unwrap_or(0);
                if failures > self.ping_failures_tolerated {
                    info!(
                        NetworkSchema::new(&self.network_context).remote_peer(&peer_id),
                        "{} Disconnecting from peer: {}",
                        self.network_context,
                        peer_id.short_str()
                    );
                    let peer_network_id =
                        PeerNetworkId::new(self.network_context.network_id(), peer_id);
                    if let Err(err) = self
                        .network_interface
                        .disconnect_peer(peer_network_id)
                        .await
                    {
                        warn!(
                            NetworkSchema::new(&self.network_context)
                                .remote_peer(&peer_id),
                            error = ?err,
                            "{} Failed to disconnect from peer: {} with error: {:?}",
                            self.network_context,
                            peer_id.short_str(),
                            err
                        );
                    }
                }
            }
        }
    }

    async fn ping_peer(
        network_context: NetworkContext,
        network_tx: HealthCheckerNetworkSender,
        peer_id: PeerId,
        round: u64,
        nonce: u32,
        ping_timeout: Duration,
    ) -> (PeerId, u64, u32, Result<Pong, RpcError>) {
        trace!(
            NetworkSchema::new(&network_context).remote_peer(&peer_id),
            round = round,
            "{} Sending Ping request to peer: {} for round: {} nonce: {}",
            network_context,
            peer_id.short_str(),
            round,
            nonce
        );
        let res_pong_msg = network_tx
            .send_rpc(peer_id, HealthCheckerMsg::Ping(Ping(nonce)), ping_timeout)
            .await
            .and_then(|msg| match msg {
                HealthCheckerMsg::Pong(res) => Ok(res),
                _ => Err(RpcError::InvalidRpcResponse),
            });
        (peer_id, round, nonce, res_pong_msg)
    }
}
<|endoftext|>


<commit_before>use rand;
use rand::Rng;
use rand::distributions::{IndependentSample, Range};
use std::thread;
use std::sync::{Mutex, MutexGuard};
use time;
use std::time::Duration;
use std::collections::HashMap;
use self::state_machine::master::MasterStateMachine;

#[macro_use]
mod state_machine;

trait RaftMsg {
    fn encode(&self) -> (usize, Vec<u8>);
}

const CHECKER_MS: u64 = 50;

//               sm , fn , data
type LogEntry = (u64, u64, Vec<u8>);
type LogEntries = Vec<LogEntry>;

service! {
    rpc append_entries(term: u64, leaderId: u64, prev_log_id: u64, prev_log_term: u64, entries: Option<LogEntries>, leader_commit: u64) -> u64; //Err for not success
    rpc request_vote(term: u64, candidate_id: u64, last_log_id: u64, last_log_term: u64) -> (u64, bool); // term, voteGranted
    rpc install_snapshot(term: u64, leader_id: u64, last_included_index: u64, last_included_term: u64, data: Vec<u8>, done: bool) -> u64;
}

fn gen_rand(lower: u64, higher: u64) -> u64 {
    let between = Range::new(lower, higher);
    let mut rng = rand::thread_rng();
    between.ind_sample(&mut rng) + 1
}

fn get_time() -> u64 {
    let timespec = time::get_time();
    let mills: f64 = timespec.sec as f64 + (timespec.nsec as f64 / 1000.0 / 1000.0 / 1000.0 );
    mills as u64
}

pub enum Membership {
    LEADER,
    FOLLOWER,
    CANDIDATE,
    OFFLINE,
}

pub struct RaftMeta {
    term: u64,
    log: u64,
    voted: bool,
    timeout: u64,
    last_checked: u64,
    last_updated: u64,
    membership: Membership,
    leader_id: u64,
    logs: LogEntries,
    state_machine: MasterStateMachine
}

#[derive(Clone)]
pub enum Storage {
    MEMORY,
    DISK(String),
}

impl Storage {
    pub fn Default() -> Storage {
        Storage::MEMORY
    }
}

#[derive(Clone)]
pub struct Options {
    pub storage: Storage,
    pub address: String,
}

pub struct RaftServer {
    meta: Arc<Mutex<RaftMeta>>,
    pub options: Options,
}

impl RaftServer {
    pub fn new(opts: Options) -> Arc<RaftServer> {
        let server = Arc::new(RaftServer {
            meta: Arc::new(Mutex::new(
                RaftMeta {
                    term: 0,
                    log: 0,
                    voted: false,
                    timeout: gen_rand(100, 500), // 10~500 ms for timeout
                    last_checked: get_time(),
                    last_updated: get_time(),
                    membership: Membership::FOLLOWER,
                    leader_id: 0,
                    logs: Vec::new(),
                    state_machine: MasterStateMachine::new(),
                }
            )),
            options: opts.clone(),
        });
        let svr_ref = server.clone();
        thread::spawn(move ||{
            listen(svr_ref, &opts.address);
        });
        let checker_ref = server.clone();
        thread::spawn(move ||{
            let server = checker_ref;
            loop {
                {
                    let mut meta = server.meta.lock().unwrap(); //WARNING: Reentering not supported
                    match meta.membership {
                        Membership::LEADER => {
                            if get_time() > (meta.last_updated + CHECKER_MS) {
                                server.send_heartbeat(&meta, None);
                            }
                        },
                        Membership::FOLLOWER | Membership::CANDIDATE => {
                            if get_time() > (meta.timeout + meta.last_checked) { //Timeout, require election
                                server.become_candidate(&meta);
                            }
                        },
                        Membership::OFFLINE => {
                            break;
                        }
                    }
                }
                thread::sleep(Duration::from_millis(CHECKER_MS));
            }
        });
        server
    }
    pub fn become_candidate(&self, meta: &MutexGuard<RaftMeta>) {

    }
    pub fn send_heartbeat(&self, meta: &MutexGuard<RaftMeta>, entries: Option<LogEntries>) {

    }
}

impl Server for RaftServer {
    fn append_entries(
        &self,
        term: u64, leaderId: u64, prev_log_id: u64,
        prev_log_term: u64, entries: Option<LogEntries>,
        leader_commit: u64
    ) -> Result<u64, ()>  {
        let mut meta = self.meta.lock().unwrap();
        Ok(meta.term)
    }

    fn request_vote(
        &self,
        term: u64, candidate_id: u64,
        last_log_id: u64, last_log_term: u64
    ) -> Result<(u64, bool), ()> {
        let mut meta = self.meta.lock().unwrap();
        let voted = meta.voted;
        if !voted {

        }
        Ok((meta.term, !voted))
    }

    fn install_snapshot(
        &self,
        term: u64, leader_id: u64, last_included_index: u64,
        last_included_term: u64, data: Vec<u8>, done: bool
    ) -> Result<u64, ()> {
        let mut meta = self.meta.lock().unwrap();
        Ok(meta.term)
    }
}
<commit_msg>basic rules<commit_after>use rand;
use rand::Rng;
use rand::distributions::{IndependentSample, Range};
use std::thread;
use std::sync::{Mutex, MutexGuard};
use time;
use std::time::Duration;
use std::collections::HashMap;
use self::state_machine::master::MasterStateMachine;

#[macro_use]
mod state_machine;

trait RaftMsg {
    fn encode(&self) -> (usize, Vec<u8>);
}

const CHECKER_MS: u64 = 50;

//               sm , fn , data
type LogEntry = (u64, u64, Vec<u8>);
type LogEntries = Vec<LogEntry>;

service! {
    rpc append_entries(term: u64, leaderId: u64, prev_log_id: u64, prev_log_term: u64, entries: Option<LogEntries>, leader_commit: u64) -> (u64, bool);
    rpc request_vote(term: u64, candidate_id: u64, last_log_id: u64, last_log_term: u64) -> (u64, bool); // term, voteGranted
    rpc install_snapshot(term: u64, leader_id: u64, last_included_index: u64, last_included_term: u64, data: Vec<u8>, done: bool) -> u64;
}

fn gen_rand(lower: u64, higher: u64) -> u64 {
    let between = Range::new(lower, higher);
    let mut rng = rand::thread_rng();
    between.ind_sample(&mut rng) + 1
}

fn get_time() -> u64 {
    let timespec = time::get_time();
    let mills: f64 = timespec.sec as f64 + (timespec.nsec as f64 / 1000.0 / 1000.0 / 1000.0 );
    mills as u64
}

pub enum Membership {
    LEADER,
    FOLLOWER,
    CANDIDATE,
    OFFLINE,
}

pub struct RaftMeta {
    term: u64,
    log: u64,
    voted: bool,
    timeout: u64,
    last_checked: u64,
    last_updated: u64,
    membership: Membership,
    leader_id: u64,
    logs: LogEntries,
    state_machine: MasterStateMachine
}

#[derive(Clone)]
pub enum Storage {
    MEMORY,
    DISK(String),
}

impl Storage {
    pub fn Default() -> Storage {
        Storage::MEMORY
    }
}

#[derive(Clone)]
pub struct Options {
    pub storage: Storage,
    pub address: String,
}

pub struct RaftServer {
    meta: Arc<Mutex<RaftMeta>>,
    pub options: Options,
}

impl RaftServer {
    pub fn new(opts: Options) -> Arc<RaftServer> {
        let server = Arc::new(RaftServer {
            meta: Arc::new(Mutex::new(
                RaftMeta {
                    term: 0,
                    log: 0,
                    voted: false,
                    timeout: gen_rand(100, 500), // 10~500 ms for timeout
                    last_checked: get_time(),
                    last_updated: get_time(),
                    membership: Membership::FOLLOWER,
                    leader_id: 0,
                    logs: Vec::new(),
                    state_machine: MasterStateMachine::new(),
                }
            )),
            options: opts.clone(),
        });
        let svr_ref = server.clone();
        thread::spawn(move ||{
            listen(svr_ref, &opts.address);
        });
        let checker_ref = server.clone();
        thread::spawn(move ||{
            let server = checker_ref;
            loop {
                {
                    let mut meta = server.meta.lock().unwrap(); //WARNING: Reentering not supported
                    match meta.membership {
                        Membership::LEADER => {
                            if get_time() > (meta.last_updated + CHECKER_MS) {
                                server.send_heartbeat(&mut meta, None);
                            }
                        },
                        Membership::FOLLOWER | Membership::CANDIDATE => {
                            if get_time() > (meta.timeout + meta.last_checked) { //Timeout, require election
                                server.become_candidate(&mut meta);
                            }
                        },
                        Membership::OFFLINE => {
                            break;
                        }
                    }
                }
                thread::sleep(Duration::from_millis(CHECKER_MS));
            }
        });
        server
    }
    fn become_candidate(&self, meta: &mut MutexGuard<RaftMeta>) {
        meta.membership = Membership::CANDIDATE;
    }
    fn become_follower(&self, meta: &mut MutexGuard<RaftMeta>) {
        meta.membership = Membership::FOLLOWER;
    }
    fn send_heartbeat(&self, meta: &mut MutexGuard<RaftMeta>, entries: Option<LogEntries>) {

    }

    //check term number, return reject = false if server term is stale
    fn check_term(&self, meta: &mut MutexGuard<RaftMeta>, remote_term: u64) -> bool {
        if remote_term > meta.term {
            (*meta).term = remote_term;
            self.become_follower(meta)
        } else if remote_term < meta.term {
            return false;
        }
        return true;
    }
}

impl Server for RaftServer {
    fn append_entries(
        &self,
        term: u64, leaderId: u64, prev_log_id: u64,
        prev_log_term: u64, entries: Option<LogEntries>,
        leader_commit: u64
    ) -> Result<(u64, bool), ()>  {
        let mut meta = self.meta.lock().unwrap();
        let term_ok = self.check_term(&mut meta, term);
        if !term_ok {
            Ok((meta.term, false))
        } else {
            Ok((meta.term, true))
        }
    }

    fn request_vote(
        &self,
        term: u64, candidate_id: u64,
        last_log_id: u64, last_log_term: u64
    ) -> Result<(u64, bool), ()> {
        let mut meta = self.meta.lock().unwrap();
        let term_ok = self.check_term(&mut meta, term);
        let voted = meta.voted;
        if term_ok && !voted {

        }
        Ok((meta.term, !voted))
    }

    fn install_snapshot(
        &self,
        term: u64, leader_id: u64, last_included_index: u64,
        last_included_term: u64, data: Vec<u8>, done: bool
    ) -> Result<u64, ()> {
        let mut meta = self.meta.lock().unwrap();
        let term_ok = self.check_term(&mut meta, term);
        Ok(meta.term)
    }
}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, WriteBytesExt};
use std::ptr;
use libc;

bitflags! {
    pub struct EntryTypes: u8 {
        const Cell = 0b00010000;
        const Tomestone = 0b00110000;
    }
}

pub struct Entry {
    entry_flag: u8,
    entry_length: u32,
}

fn count_len_bytes(len: u32) -> u8 {
    let in_bits = 32;
    let msb = 1 << (in_bits - 1);
    let mut count: u8 = 0;
    for i in 0..in_bits
    {
        if (len << i) & msb > 0 {
            break;
        };
        count += 1;
    }
    let bytes = count / 8;
    assert!(bytes <= 4);
    return bytes;
}

fn compress_len(len: u32, bytes: &mut[u8]) {
    bytes.write_u32::<LittleEndian>(len).unwrap();
}

impl Entry {
    pub fn encode_to<W>(mut pos: usize, entry_type: EntryTypes, content_len: u32, write_content: W)
        where W: Fn(usize)
    {
        let len_bytes_count = count_len_bytes(content_len);
        let len_bytes_count_usize = len_bytes_count as usize;
        let flag_byte = len_bytes_count | entry_type.bits;
        let mut len_bytes = Vec::with_capacity(len_bytes_count_usize);
        compress_len(content_len, &mut len_bytes);
        let raw_len_bytes= Box::into_raw(len_bytes.into_boxed_slice());
        unsafe {
            // write entry flags
            *(pos as *mut u8) = flag_byte;
            pos += 1;
            // write entry length
            libc::memmove(
                pos as *mut libc::c_void,
                raw_len_bytes as *mut libc::c_void,
                len_bytes_count_usize);
            pos += len_bytes_count_usize;
            write_content(pos);
            // release raw pointers
            Box::from_raw(raw_len_bytes);
        }
    }

    pub fn decode_from(mut pos: usize) -> (Entry, usize) {
        unsafe {
            let flag_byte = *(pos as *mut u8);
            let entry_type_bits = 0b11110000 & flag_byte;
            let entry_type = EntryTypes::from_bits(entry_type_bits);
            let entry_bytes_len = 0b00001111 & flag_byte;

        }
    }
}

<commit_msg>impl decoder<commit_after>use byteorder::{LittleEndian, WriteBytesExt, ByteOrder};
use std::ptr;
use libc;

bitflags! {
    pub struct EntryType: u8 {
        const Cell = 0b00010000;
        const Tomestone = 0b00110000;
    }
}

pub struct EntryHeader {
    entry_type: EntryType,
    entry_length: u32,
}

fn count_len_bytes(len: u32) -> u8 {
    let in_bits = 32;
    let msb = 1 << (in_bits - 1);
    let mut count: u8 = 0;
    for i in 0..in_bits
    {
        if (len << i) & msb > 0 {
            break;
        };
        count += 1;
    }
    let bytes = count / 8;
    assert!(bytes <= 4);
    return bytes;
}

fn encode_len(len: u32, bytes: &mut[u8]) {
    LittleEndian::write_u32(bytes, len);
}

impl EntryHeader {
    pub fn encode_to<W>(mut pos: usize, entry_type: EntryType, content_len: u32, write_content: W)
        where W: Fn(usize)
    {
        let len_bytes_count = count_len_bytes(content_len);
        let len_bytes_count_usize = len_bytes_count as usize;
        let flag_byte = len_bytes_count | entry_type.bits;
        let mut len_bytes = [0u8; 4];
        encode_len(content_len, &mut len_bytes);
        let raw_len_bytes= Box::into_raw(box len_bytes);
        unsafe {
            // write entry flags
            *(pos as *mut u8) = flag_byte;
            pos += 1;
            // write entry length
            libc::memmove(
                pos as *mut libc::c_void,
                raw_len_bytes as *mut libc::c_void,
                len_bytes_count_usize);
            pos += len_bytes_count_usize;
            write_content(pos);
            // release raw pointers
            Box::from_raw(raw_len_bytes);
        }
    }

    // Returns the entry header and content position
    pub fn decode_from(mut pos: usize) -> (EntryHeader, usize) {
        unsafe {
            let flag_byte = *(pos as *mut u8);
            pos += 1;
            let entry_type_bits = 0b11110000 & flag_byte;
            let entry_type = EntryType::from_bits(entry_type_bits).unwrap();
            let entry_bytes_len = 0b00001111 & flag_byte;
            let entry_bytes_len_usize = entry_bytes_len as usize;
            let raw_len_bytes= Box::into_raw(box [0u8; 4]);
            libc::memmove(
                pos as *mut libc::c_void,
                raw_len_bytes as *mut libc::c_void,
                entry_bytes_len_usize);
            let entry_length = LittleEndian::read_u32(&*Box::from_raw(raw_len_bytes));
            let enrty = EntryHeader {
                entry_type,
                entry_length
            };
            (entry, pos + 1 + entry_bytes_len_usize)
        }
    }
}

<|endoftext|>


<commit_before>use parser::*;

/// Renders Textile string into HTML string.
///
/// # Example
///
/// ```rust
/// let html = textile::render("h2. *Header of level 2*");
/// assert_eq!(html, "<h2><strong>Header of level 2</strong></h2>".to_string());
/// ```
pub fn render(text: &str) -> String {
    render_blocks(parse(text))
}

fn render_attributes(attributes: &Attributes) -> String {
     if !attributes.is_empty() {
         let mut attrs = Vec::new();

         for (key, value) in attributes {
             attrs.push(format!("{}: {}", key, value));
         }

         format!(" style=\"{};\"", attrs.join("; "))
     } else {
         String::default()
     }
}

fn render_blocks(elements: Vec<Block>) -> String {
    let mut res = String::new();

    for element in &elements {
        let html = match *element {
            Block::Header {ref attributes, level, ref elements} => format!("<h{0}{1}>{2}</h{0}>", level, render_attributes(attributes), render_inline_elements(elements)),
            Block::Paragraph {ref elements} => format!("
{}
", render_inline_elements(elements)),
            Block::BlockQuotation(ref elements) => format!("
{}
", render_inline_elements(elements)),
            Block::Code(ref code) => format!("
{}
", code)
        };
        res.push_str(&html);
    }
    res
}

fn render_inline_elements(elements: &Vec<Inline>) -> String {
    let mut res = String::new();

    for element in elements {
        let html = match *element {
            Inline::Break => "
".to_string(),
            Inline::Text(ref text) => text.to_string(),
            Inline::Bold(ref elements, ref tag_type) => {
                let tag = match *tag_type {
                    BoldTagType::Strong => "strong",
                    BoldTagType::Bold => "b",
                };
                format!("<{0}>{1}</{0}>", tag, render_inline_elements(elements))
            },
            Inline::Italic(ref elements, ref tag_type) => {
                let tag = match *tag_type {
                    ItalicTagType::Emphasis => "em",
                    ItalicTagType::Italic => "i",
                };
                format!("<{0}>{1}</{0}>", tag, render_inline_elements(elements))
            },
            Inline::Strikethrough(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Underlined(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Subscript(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Superscript(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Code(ref text) => format!("{}", text),
            Inline::Citation(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Abbreviation {ref abbr, ref transcript} => format!("<acronym title=\"{}\">{}</acronym>", transcript, abbr),
            Inline::Link {ref description, ref url} => format!("{}", url, render_inline_elements(description)),
            Inline::Image {ref url, alt: Some(ref text)} => format!("", url, text),
            Inline::Image {ref url, alt: None} => format!("", url),
        };
        res.push_str(&html);
    }
    res
}
<commit_msg>Change argument type in render_inline_elements function<commit_after>use parser::*;

/// Renders Textile string into HTML string.
///
/// # Example
///
/// ```rust
/// let html = textile::render("h2. *Header of level 2*");
/// assert_eq!(html, "<h2><strong>Header of level 2</strong></h2>".to_string());
/// ```
pub fn render(text: &str) -> String {
    render_blocks(parse(text))
}

fn render_attributes(attributes: &Attributes) -> String {
     if !attributes.is_empty() {
         let mut attrs = Vec::new();

         for (key, value) in attributes {
             attrs.push(format!("{}: {}", key, value));
         }

         format!(" style=\"{};\"", attrs.join("; "))
     } else {
         String::default()
     }
}

fn render_blocks(elements: Vec<Block>) -> String {
    let mut res = String::new();

    for element in &elements {
        let html = match *element {
            Block::Header {ref attributes, level, ref elements} => format!("<h{0}{1}>{2}</h{0}>", level, render_attributes(attributes), render_inline_elements(elements)),
            Block::Paragraph {ref elements} => format!("
{}
", render_inline_elements(elements)),
            Block::BlockQuotation(ref elements) => format!("
{}
", render_inline_elements(elements)),
            Block::Code(ref code) => format!("
{}
", code)
        };
        res.push_str(&html);
    }
    res
}

fn render_inline_elements(elements: &[Inline]) -> String {
    let mut res = String::new();

    for element in elements {
        let html = match *element {
            Inline::Break => "
".to_string(),
            Inline::Text(ref text) => text.to_string(),
            Inline::Bold(ref elements, ref tag_type) => {
                let tag = match *tag_type {
                    BoldTagType::Strong => "strong",
                    BoldTagType::Bold => "b",
                };
                format!("<{0}>{1}</{0}>", tag, render_inline_elements(elements))
            },
            Inline::Italic(ref elements, ref tag_type) => {
                let tag = match *tag_type {
                    ItalicTagType::Emphasis => "em",
                    ItalicTagType::Italic => "i",
                };
                format!("<{0}>{1}</{0}>", tag, render_inline_elements(elements))
            },
            Inline::Strikethrough(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Underlined(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Subscript(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Superscript(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Code(ref text) => format!("{}", text),
            Inline::Citation(ref elements) => format!("{}", render_inline_elements(elements)),
            Inline::Abbreviation {ref abbr, ref transcript} => format!("<acronym title=\"{}\">{}</acronym>", transcript, abbr),
            Inline::Link {ref description, ref url} => format!("{}", url, render_inline_elements(description)),
            Inline::Image {ref url, alt: Some(ref text)} => format!("", url, text),
            Inline::Image {ref url, alt: None} => format!("", url),
        };
        res.push_str(&html);
    }
    res
}
<|endoftext|>


<commit_before>use serde::de::{self, Deserializer, Visitor};
use serde::Deserialize;
use std::fmt;
use std::str::FromStr;
use std::time::SystemTime;

#[derive(Debug, Deserialize, PartialEq)]
#[serde(rename = "usage_report")]
pub struct UsageReport {
    pub metric: String,
    pub period: Period,
    pub period_start: String,
    pub period_end: String,
    pub max_value: u64,
    pub current_value: u64,
}

// Unfortunately the XML output from Apisonator includes a rather useless "usage_reports" tag that
// is then followed by a "usage_report" tag in each UsageReport, so we need to wrap that up.
#[derive(Debug, Deserialize, PartialEq)]
pub enum UsageReports {
    #[serde(rename = "usage_report")]
    UsageReports(Vec<UsageReport>),
}

#[derive(Debug, Deserialize, PartialEq)]
#[serde(rename = "status")]
pub struct Authorization {
    authorized: bool,
    plan: String,
    usage_reports: UsageReports,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct PeriodInstance {
    start: SystemTime,
    end: SystemTime,
}

#[derive(Debug, PartialEq)]
pub enum Period {
    Minute,
    Hour,
    Day,
    Week,
    Month,
    Year,
    Eternity,
}

struct PeriodStringVisitor;

impl<'de> Visitor<'de> for PeriodStringVisitor {
    type Value = Period;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a string that represents a period")
    }

    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        match v {
            "minute" => Ok(Period::Minute),
            "hour" => Ok(Period::Hour),
            "day" => Ok(Period::Day),
            "week" => Ok(Period::Week),
            "month" => Ok(Period::Month),
            "year" => Ok(Period::Year),
            "eternity" => Ok(Period::Eternity),
            _ => Err(E::custom("Invalid period")),
        }
    }
}

impl<'de> Deserialize<'de> for Period {
    fn deserialize<D>(deserializer: D) -> Result<Period, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_any(PeriodStringVisitor)
    }
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct UsageData {
    max_value: u64,
    current_value: u64,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct Metric(pub String);

impl FromStr for Authorization {
    type Err = serde_xml_rs::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        serde_xml_rs::from_str(s)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse() {
        let s = r##"
        <?xml version="1.0" encoding="UTF-8"?>
        <status>
            <authorized>true</authorized>
            <plan>App Plan</plan>
            <usage_reports>
                <usage_report metric="products" period="minute">
                    <period_start>2019-06-05 16:24:00 +0000</period_start>
                    <period_end>2019-06-05 16:25:00 +0000</period_end>
                    <max_value>5</max_value>
                    <current_value>0</current_value>
                </usage_report>
                <usage_report metric="products" period="month">
                    <period_start>2019-06-01 00:00:00 +0000</period_start>
                    <period_end>2019-07-01 00:00:00 +0000</period_end>
                    <max_value>50</max_value>
                    <current_value>0</current_value>
                </usage_report>
            </usage_reports>
        </status>
        "##;

        let parsed_auth = Authorization::from_str(s).unwrap();

        let expected_auth = Authorization {
            authorized: true,
            plan: String::from("App Plan"),
            usage_reports: UsageReports::UsageReports(vec![
                UsageReport {
                    metric: String::from("products"),
                    period: Period::Minute,
                    period_start: String::from("2019-06-05 16:24:00 +0000"),
                    period_end: String::from("2019-06-05 16:25:00 +0000"),
                    max_value: 5,
                    current_value: 0,
                },
                UsageReport {
                    metric: String::from("products"),
                    period: Period::Month,
                    period_start: String::from("2019-06-01 00:00:00 +0000"),
                    period_end: String::from("2019-07-01 00:00:00 +0000"),
                    max_value: 50,
                    current_value: 0,
                },
            ]),
        };

        assert_eq!(parsed_auth, expected_auth);
    }
}
<commit_msg>response: expose periods as SystemTime instances instead of strings<commit_after>use chrono::prelude::*;
use serde::de::{self, Deserializer, MapAccess, Visitor};
use serde::Deserialize;
use std::fmt;
use std::str::FromStr;
use std::time::SystemTime;

#[derive(Debug, PartialEq)]
pub struct PeriodTime(SystemTime);

#[derive(Debug, Deserialize, PartialEq)]
#[serde(rename = "usage_report")]
pub struct UsageReport {
    pub metric: String,
    pub period: Period,
    pub period_start: PeriodTime,
    pub period_end: PeriodTime,
    pub max_value: u64,
    pub current_value: u64,
}

// Unfortunately the XML output from Apisonator includes a rather useless "usage_reports" tag that
// is then followed by a "usage_report" tag in each UsageReport, so we need to wrap that up.
#[derive(Debug, Deserialize, PartialEq)]
pub enum UsageReports {
    #[serde(rename = "usage_report")]
    UsageReports(Vec<UsageReport>),
}

#[derive(Debug, Deserialize, PartialEq)]
#[serde(rename = "status")]
pub struct Authorization {
    authorized: bool,
    plan: String,
    usage_reports: UsageReports,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct PeriodInstance {
    start: SystemTime,
    end: SystemTime,
}

#[derive(Debug, PartialEq)]
pub enum Period {
    Minute,
    Hour,
    Day,
    Week,
    Month,
    Year,
    Eternity,
}

struct PeriodStringVisitor;

impl<'de> Visitor<'de> for PeriodStringVisitor {
    type Value = Period;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a string that represents a period")
    }

    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        match v {
            "minute" => Ok(Period::Minute),
            "hour" => Ok(Period::Hour),
            "day" => Ok(Period::Day),
            "week" => Ok(Period::Week),
            "month" => Ok(Period::Month),
            "year" => Ok(Period::Year),
            "eternity" => Ok(Period::Eternity),
            _ => Err(E::custom("Invalid period")),
        }
    }
}

impl<'de> Deserialize<'de> for Period {
    fn deserialize<D>(deserializer: D) -> Result<Period, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_any(PeriodStringVisitor)
    }
}

struct TimestampVisitor;

impl<'de> Visitor<'de> for TimestampVisitor {
    type Value = PeriodTime;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a string that represents a timestamp")
    }

    fn visit_map<V>(self, mut map: V) -> Result<PeriodTime, V::Error>
    where
        V: MapAccess<'de>,
    {
        // We know there's only one key with one value - Note: probably safe to skip
        let _key: Option<String> = map.next_key()?;
        let timestamp: String = map.next_value()?;

        let dt = DateTime::parse_from_str(timestamp.as_str(), "%Y-%m-%d %H:%M:%S %z").unwrap();

        Ok(PeriodTime(dt.into()))
    }
}

impl<'de> Deserialize<'de> for PeriodTime {
    fn deserialize<D>(deserializer: D) -> Result<PeriodTime, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_any(TimestampVisitor)
    }
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct UsageData {
    max_value: u64,
    current_value: u64,
}

#[derive(Debug, Deserialize, PartialEq)]
pub struct Metric(pub String);

impl FromStr for Authorization {
    type Err = serde_xml_rs::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        serde_xml_rs::from_str(s)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse() {
        let s = r##"
        <?xml version="1.0" encoding="UTF-8"?>
        <status>
            <authorized>true</authorized>
            <plan>App Plan</plan>
            <usage_reports>
                <usage_report metric="products" period="minute">
                    <period_start>2019-06-05 16:24:00 +0000</period_start>
                    <period_end>2019-06-05 16:25:00 +0000</period_end>
                    <max_value>5</max_value>
                    <current_value>0</current_value>
                </usage_report>
                <usage_report metric="products" period="month">
                    <period_start>2019-06-01 00:00:00 +0000</period_start>
                    <period_end>2019-07-01 00:00:00 +0000</period_end>
                    <max_value>50</max_value>
                    <current_value>0</current_value>
                </usage_report>
            </usage_reports>
        </status>
        "##;

        let parsed_auth = Authorization::from_str(s).unwrap();

        let expected_auth = Authorization {
            authorized: true,
            plan: String::from("App Plan"),
            usage_reports: UsageReports::UsageReports(vec![
                UsageReport {
                    metric: String::from("products"),
                    period: Period::Minute,
                    period_start: PeriodTime(Utc.ymd(2019, 6, 5).and_hms(16, 24, 0).into()),
                    period_end: PeriodTime(Utc.ymd(2019, 6, 5).and_hms(16, 25, 0).into()),
                    max_value: 5,
                    current_value: 0,
                },
                UsageReport {
                    metric: String::from("products"),
                    period: Period::Month,
                    period_start: PeriodTime(Utc.ymd(2019, 6, 1).and_hms(0, 0, 0).into()),
                    period_end: PeriodTime(Utc.ymd(2019, 7, 1).and_hms(0, 0, 0).into()),
                    max_value: 50,
                    current_value: 0,
                },
            ]),
        };

        assert_eq!(parsed_auth, expected_auth);
    }
}
<|endoftext|>


<commit_before>//! Types for the *m.room_key* event.

use ruma_events_macros::ruma_event;
use ruma_identifiers::RoomId;

use super::Algorithm;

ruma_event! {
    /// This event type is used to exchange keys for end-to-end encryption.
    ///
    /// Typically it is encrypted as an *m.room.encrypted* event, then sent as a to-device event.
    RoomKeyEvent {
        kind: Event,
        event_type: "m.room_key",
        content: {
            /// The encryption algorithm the key in this event is to be used with.
            ///
            /// Must be `m.megolm.v1.aes-sha2`.
            pub algorithm: Algorithm,

            /// The room where the key is used.
            pub room_id: RoomId,

            /// The ID of the session that the key is for.
            pub session_id: String,

            /// The key to be exchanged.
            pub session_key: String,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::convert::TryFrom;

    use ruma_identifiers::RoomId;
    use serde_json::{json, to_value as to_json_value};

    use super::{RoomKeyEvent, RoomKeyEventContent};
    use crate::Algorithm;

    #[test]
    fn serialization() {
        let ev = RoomKeyEvent {
            content: RoomKeyEventContent {
                algorithm: Algorithm::MegolmV1AesSha2,
                room_id: RoomId::try_from("!testroomid:exmaple.org").unwrap(),
                session_id: "SessId".into(),
                session_key: "SessKey".into(),
            },
        };

        assert_eq!(
            to_json_value(ev).unwrap(),
            json!({
                "type": "m.room_key",
                "content": {
                    "algorithm": "m.megolm.v1.aes-sha2",
                    "room_id": "!testroomid:exmaple.org",
                    "session_id": "SessId",
                    "session_key": "SessKey",
                },
            })
        );
    }
}
<commit_msg>Fix typo<commit_after>//! Types for the *m.room_key* event.

use ruma_events_macros::ruma_event;
use ruma_identifiers::RoomId;

use super::Algorithm;

ruma_event! {
    /// This event type is used to exchange keys for end-to-end encryption.
    ///
    /// Typically it is encrypted as an *m.room.encrypted* event, then sent as a to-device event.
    RoomKeyEvent {
        kind: Event,
        event_type: "m.room_key",
        content: {
            /// The encryption algorithm the key in this event is to be used with.
            ///
            /// Must be `m.megolm.v1.aes-sha2`.
            pub algorithm: Algorithm,

            /// The room where the key is used.
            pub room_id: RoomId,

            /// The ID of the session that the key is for.
            pub session_id: String,

            /// The key to be exchanged.
            pub session_key: String,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::convert::TryFrom;

    use ruma_identifiers::RoomId;
    use serde_json::{json, to_value as to_json_value};

    use super::{RoomKeyEvent, RoomKeyEventContent};
    use crate::Algorithm;

    #[test]
    fn serialization() {
        let ev = RoomKeyEvent {
            content: RoomKeyEventContent {
                algorithm: Algorithm::MegolmV1AesSha2,
                room_id: RoomId::try_from("!testroomid:example.org").unwrap(),
                session_id: "SessId".into(),
                session_key: "SessKey".into(),
            },
        };

        assert_eq!(
            to_json_value(ev).unwrap(),
            json!({
                "type": "m.room_key",
                "content": {
                    "algorithm": "m.megolm.v1.aes-sha2",
                    "room_id": "!testroomid:example.org",
                    "session_id": "SessId",
                    "session_key": "SessKey",
                },
            })
        );
    }
}
<|endoftext|>


<commit_before>use symbols::{file_name, symbol_name};
use expressions::{expression, value, Expression};
use statements::{statement, Statement};

#[derive(Debug, PartialEq)]
pub enum InputSection {
    Section(String),
    ExcludeFile(String),
    Sort(String),
}

#[derive(Debug, PartialEq)]
pub enum SectionItem {
    Statement(Statement),
    Sections {
        file: String,
        sections: Vec<InputSection>,
    },
    File(String),
    Keep(Box<SectionItem>),
}

#[derive(Debug, PartialEq)]
pub struct OutputSection {
    pub name: String,
    pub start: Option<Expression>,
    pub align: Option<Expression>,
    pub no_load: bool,
    pub load_address: Option<Expression>,
    pub contents: Vec<SectionItem>,
    pub region: Option<String>,
    pub region_at: Option<String>,
    pub fill: Option<Expression>,
}

#[derive(Debug, PartialEq)]
pub enum Section {
    Definition(OutputSection),
    Statement(Statement),
}

named!(input_section_name<&str, InputSection>, map!(
    file_name,
    |x: &str| InputSection::Section(x.into())
));

named!(input_section_exclude<&str, InputSection>, wsc!(do_parse!(
    tag_s!("EXCLUDE_FILE")
    >>
    tag_s!("(")
    >>
    name: file_name
    >>
    tag_s!(")")
    >>
    (InputSection::ExcludeFile(name.into()))
)));

named!(input_section_sort<&str, InputSection>, wsc!(do_parse!(
    tag_s!("SORT")
    >>
    tag_s!("(")
    >>
    name: file_name
    >>
    tag_s!(")")
    >>
    (InputSection::Sort(name.into()))
)));

named!(input_section<&str, InputSection>, alt_complete!(
    input_section_sort | input_section_exclude | input_section_name
));

named!(section_item_sections<&str, SectionItem>, wsc!(do_parse!(
    name: file_name
    >>
    tag_s!("(")
    >>
    sections: many1!(input_section)
    >>
    tag_s!(")")
    >>
    opt!(tag_s!(";"))
    >>
    (SectionItem::Sections{
        file: name.into(),
        sections: sections,
    })
)));

named!(section_item_file<&str, SectionItem>, map!(
    file_name,
    |x: &str| SectionItem::File(x.into())
));

named!(section_item_keep<&str, SectionItem>, wsc!(do_parse!(
    tag_s!("KEEP")
    >>
    tag_s!("(")
    >>
    item: section_item
    >>
    tag_s!(")")
    >>
    opt!(tag_s!(";"))
    >>
    (SectionItem::Keep(Box::new(item)))
)));

named!(section_item_statement<&str, SectionItem>, map!(
    statement,
    |x| SectionItem::Statement(x)
));

named!(section_item<&str, SectionItem>, alt_complete!(
    section_item_keep | section_item_sections | section_item_statement | section_item_file
));

named!(section_items<&str, Vec<SectionItem>>, many0!(
    section_item
));

named!(align<&str, Expression>, wsc!(do_parse!(
    tag_s!("BLOCK")
    >>
    tag_s!("(")
    >>
    expr: expression
    >>
    tag_s!(")")
    >>
    (expr)
)));

named!(load_addr<&str, Expression>, wsc!(do_parse!(
    tag_s!("AT")
    >>
    tag_s!("(")
    >>
    expr: expression
    >>
    tag_s!(")")
    >>
    (expr)
)));

named!(region<&str, String>, wsc!(do_parse!(
    tag_s!(">")
    >>
    name: symbol_name
    >>
    (name.into())
)));

named!(region_at<&str, String>, wsc!(do_parse!(
    tag_s!("AT")
    >>
    tag_s!(">")
    >>
    name: symbol_name
    >>
    (name.into())
)));

named!(fill<&str, Expression>, wsc!(do_parse!(
    tag_s!("=")
    >>
    expr: value
    >>
    (expr)
)));

named!(output_section<&str, OutputSection>, wsc!(do_parse!(
    sect_name: alt_complete!(tag_s!("/DISCARD/") | symbol_name)
    >>
    start: opt!(value)
    >>
    align: opt!(align)
    >>
    no_load: opt!(tag_s!("(NOLOAD)"))
    >>
    tag_s!(":")
    >>
    load_address: opt!(load_addr)
    >>
    tag_s!("{")
    >>
    items: section_items
    >>
    tag_s!("}")
    >>
    region: opt!(region)
    >>
    region_at: opt!(region_at)
    >>
    fill: opt!(fill)
    >>
    (OutputSection{
        name: sect_name.into(),
        start: start,
        align: align,
        no_load: no_load.is_some(),
        load_address: load_address,
        contents: items,
        region: region,
        region_at: region_at,
        fill: fill,
    })
)));

named!(sect_statement<&str, Section>, map!(
    statement,
    |x| Section::Statement(x)
));

named!(sect_definition<&str, Section>, map!(
    output_section,
    |x| Section::Definition(x)
));

named!(pub sections<&str, Vec<Section>>, many0!(alt_complete!(
    sect_definition | sect_statement
)));

#[cfg(test)]
mod test {
    use nom::IResult;
    use statements::Statement;
    use expressions::Expression::Number;
    use sections::{section_items, SectionItem, InputSection};

    #[test]
    fn test_sections() {
        match section_items(" a ( b* .g ) ; KEEP ( * ( SORT ( c ) ) ) foo.o . = 0 ; ") {
            IResult::Done("", v) => {
                assert_eq!(v.len(), 4);

                assert_eq!(v[0],
                           SectionItem::Sections {
                               file: String::from("a"),
                               sections: vec![InputSection::Section(String::from("b*")),
                                              InputSection::Section(String::from(".g"))],
                           });

                assert_eq!(v[1], SectionItem::Keep(Box::new(
                    SectionItem::Sections{
                    file: String::from("*"),
                    sections: vec![InputSection::Sort(String::from("c"))]
                })));

                assert_eq!(v[2], SectionItem::File(String::from("foo.o")));

                assert_eq!(v[3],
                           SectionItem::Statement(Statement::Assign {
                                                      symbol: String::from("."),
                                                      operator: String::from("="),
                                                      expr: Number(0),
                                                  }));
            }
            _ => assert!(false),
        }
    }
}
<commit_msg>allow any command in section declarations<commit_after>use symbols::{file_name, symbol_name};
use expressions::{expression, value, Expression};
use statements::{statement, Statement};

#[derive(Debug, PartialEq)]
pub enum InputSection {
    Section(String),
    Command {
        name: String,
        arg: String
    }
}

#[derive(Debug, PartialEq)]
pub enum SectionItem {
    Statement(Statement),
    Sections {
        file: String,
        sections: Vec<InputSection>,
    },
    File(String),
    Keep(Box<SectionItem>),
}

#[derive(Debug, PartialEq)]
pub struct OutputSection {
    pub name: String,
    pub start: Option<Expression>,
    pub align: Option<Expression>,
    pub no_load: bool,
    pub load_address: Option<Expression>,
    pub contents: Vec<SectionItem>,
    pub region: Option<String>,
    pub region_at: Option<String>,
    pub fill: Option<Expression>,
}

#[derive(Debug, PartialEq)]
pub enum Section {
    Definition(OutputSection),
    Statement(Statement),
}

named!(input_section_name<&str, InputSection>, map!(
    file_name,
    |x: &str| InputSection::Section(x.into())
));

named!(input_section_command<&str, InputSection>, wsc!(do_parse!(
    name: symbol_name
    >>
    tag_s!("(")
    >>
    arg: file_name
    >>
    tag_s!(")")
    >>
    (InputSection::Command{
        name: name.into(),
        arg: arg.into()
    })
)));

named!(input_section<&str, InputSection>, alt_complete!(
    input_section_command | input_section_name
));

named!(section_item_sections<&str, SectionItem>, wsc!(do_parse!(
    name: file_name
    >>
    tag_s!("(")
    >>
    sections: many1!(input_section)
    >>
    tag_s!(")")
    >>
    opt!(tag_s!(";"))
    >>
    (SectionItem::Sections{
        file: name.into(),
        sections: sections,
    })
)));

named!(section_item_file<&str, SectionItem>, map!(
    file_name,
    |x: &str| SectionItem::File(x.into())
));

named!(section_item_keep<&str, SectionItem>, wsc!(do_parse!(
    tag_s!("KEEP")
    >>
    tag_s!("(")
    >>
    item: section_item
    >>
    tag_s!(")")
    >>
    opt!(tag_s!(";"))
    >>
    (SectionItem::Keep(Box::new(item)))
)));

named!(section_item_statement<&str, SectionItem>, map!(
    statement,
    |x| SectionItem::Statement(x)
));

named!(section_item<&str, SectionItem>, alt_complete!(
    section_item_keep | section_item_sections | section_item_statement | section_item_file
));

named!(section_items<&str, Vec<SectionItem>>, many0!(
    section_item
));

named!(align<&str, Expression>, wsc!(do_parse!(
    tag_s!("BLOCK")
    >>
    tag_s!("(")
    >>
    expr: expression
    >>
    tag_s!(")")
    >>
    (expr)
)));

named!(load_addr<&str, Expression>, wsc!(do_parse!(
    tag_s!("AT")
    >>
    tag_s!("(")
    >>
    expr: expression
    >>
    tag_s!(")")
    >>
    (expr)
)));

named!(region<&str, String>, wsc!(do_parse!(
    tag_s!(">")
    >>
    name: symbol_name
    >>
    (name.into())
)));

named!(region_at<&str, String>, wsc!(do_parse!(
    tag_s!("AT")
    >>
    tag_s!(">")
    >>
    name: symbol_name
    >>
    (name.into())
)));

named!(fill<&str, Expression>, wsc!(do_parse!(
    tag_s!("=")
    >>
    expr: value
    >>
    (expr)
)));

named!(output_section<&str, OutputSection>, wsc!(do_parse!(
    sect_name: alt_complete!(tag_s!("/DISCARD/") | symbol_name)
    >>
    start: opt!(value)
    >>
    align: opt!(align)
    >>
    no_load: opt!(tag_s!("(NOLOAD)"))
    >>
    tag_s!(":")
    >>
    load_address: opt!(load_addr)
    >>
    tag_s!("{")
    >>
    items: section_items
    >>
    tag_s!("}")
    >>
    region: opt!(region)
    >>
    region_at: opt!(region_at)
    >>
    fill: opt!(fill)
    >>
    (OutputSection{
        name: sect_name.into(),
        start: start,
        align: align,
        no_load: no_load.is_some(),
        load_address: load_address,
        contents: items,
        region: region,
        region_at: region_at,
        fill: fill,
    })
)));

named!(sect_statement<&str, Section>, map!(
    statement,
    |x| Section::Statement(x)
));

named!(sect_definition<&str, Section>, map!(
    output_section,
    |x| Section::Definition(x)
));

named!(pub sections<&str, Vec<Section>>, many0!(alt_complete!(
    sect_definition | sect_statement
)));

#[cfg(test)]
mod test {
    use nom::IResult;
    use statements::Statement;
    use expressions::Expression::Number;
    use sections::{section_items, SectionItem, InputSection};

    #[test]
    fn test_sections() {
        match section_items(" a ( b* .g ) ; KEEP ( * ( SORT ( c ) ) ) foo.o . = 0 ; ") {
            IResult::Done("", v) => {
                assert_eq!(v.len(), 4);

                assert_eq!(v[0],
                           SectionItem::Sections {
                               file: String::from("a"),
                               sections: vec![InputSection::Section(String::from("b*")),
                                              InputSection::Section(String::from(".g"))],
                           });

                assert_eq!(v[1], SectionItem::Keep(Box::new(
                    SectionItem::Sections{
                    file: String::from("*"),
                    sections: vec![InputSection::Command{
                        name: String::from("SORT"),
                        arg: String::from("c")
                    }]
                })));

                assert_eq!(v[2], SectionItem::File(String::from("foo.o")));

                assert_eq!(v[3],
                           SectionItem::Statement(Statement::Assign {
                                                      symbol: String::from("."),
                                                      operator: String::from("="),
                                                      expr: Number(0),
                                                  }));
            }
            _ => assert!(false),
        }
    }
}
<|endoftext|>


<commit_before>//challenge 9
pub fn pkcs_7(buf: &mut Vec<u8>, block_len: usize) {
    assert!(block_len < 256);
    let mut remain = block_len as isize - buf.len() as isize;
    let mut new_len = block_len;
    if remain < 0 {
        panic!("buf is larger than block size");
    } else if remain == 0  {
        new_len = block_len + block_len;
        remain = block_len as isize;
    }

    let to_pad = remain as u8;
    unsafe {
        let orig_len = buf.len();
        buf.set_len(new_len);
        let b_ptr = buf.as_mut_ptr().offset(orig_len as isize);
        for x in 0..remain {
            *b_ptr.offset(x as isize) = to_pad;
        }
    }
}

pub fn pkcs_7_remove(buf: &mut Vec<u8>) -> bool {
    let last_byte = buf[buf.len() - 1];
    if last_byte as usize > buf.len() {
        return false;
    } else {
        let orig_len = buf.len();
        buf.resize(orig_len - last_byte as usize, 0);
        return true;
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_pkcs_7_not_multiple() {
        let mut block = b"Minkowski".to_vec();
        pkcs_7(&mut block, 16);
        assert_eq!(&block, b"Minkowski\x07\x07\x07\x07\x07\x07\x07");
    }

    #[test]
    fn test_pkcs_7_multiple() {
        let mut block = b"Euclid".to_vec();
        pkcs_7(&mut block, 6);
        assert_eq!(&block, b"Euclid\x06\x06\x06\x06\x06\x06");
    }
}
<commit_msg>challenge12<commit_after>//challenge 9
pub fn pkcs_7(buf: &mut Vec<u8>, block_len: usize) {
    assert!(block_len < 256);
    let mut remain = block_len as isize - buf.len() as isize;
    let mut new_len = block_len;
    if remain < 0 {
        panic!("buf is larger than block size");
    } else if remain == 0  {
        new_len = block_len + block_len;
        remain = block_len as isize;
    }

    let to_pad = remain as u8;
    unsafe {
        let orig_len = buf.len();
        buf.set_len(new_len);
        let b_ptr = buf.as_mut_ptr().offset(orig_len as isize);
        for x in 0..remain {
            *b_ptr.offset(x as isize) = to_pad;
        }
    }
}

pub fn pkcs_7_remove(buf: &mut Vec<u8>) -> bool {
    let last_byte = buf[buf.len() - 1];
    if last_byte as usize > buf.len() {
        return false;
    } else {
        let orig_len = buf.len();
        buf.resize(orig_len - last_byte as usize, 0);
        return true;
    }
}

pub fn ecb_decypt_appended<F>(mut f: F) -> Vec<u8>
        where F: FnMut(&[u8]) -> Vec<u8> {
    let p0_len = f(&[0u8; 0]).len();

    let mut block_len = 0;
    let mut suffix_len = 0;
    for x in 1usize.. {
        let l = f(&vec![0u8; x]).len();
        if l != p0_len {
            suffix_len = p0_len - x;
            block_len = l - p0_len;
            break;
        }
    }

    let mut attacker: Vec<u8> = vec![0; suffix_len - (suffix_len % block_len) + block_len];
    let attacker_len = attacker.len();
    let mut known_suffix: Vec<u8> = vec![];

    while known_suffix.len() != suffix_len {
        let mut ans = f(&vec![0; attacker_len - known_suffix.len() - 1]);
        ans.resize(attacker_len, 0);
        attacker[attacker_len-known_suffix.len()-1 .. attacker_len-1].clone_from_slice(&known_suffix);

        for x in 0u16..256u16 {
            attacker[attacker_len-1] = x as u8;
            // ans and trial result should only differs in 1 block, so it is a waste to compare
            // the whole array, but performance is not too important here.
            if f(&attacker)[..attacker_len] == ans[..attacker_len] {
                known_suffix.push(x as u8);
                break;
            }
        }
    }

    return known_suffix;
}

#[cfg(test)]
mod test {
    use super::*;
    extern crate base64;
    extern crate openssl;
    use self::openssl::symm::{self, Cipher};

    struct MyEncrypter {
        suffix: Vec<u8>,
        key: Vec<u8>
    }

    impl MyEncrypter {
        pub fn new() -> Self {
            let secret = b"Um9sbGluJyBpbiBteSA1LjAKV2l0aCBteSByYWctdG9wIGRvd24gc28gbXkg
aGFpciBjYW4gYmxvdwpUaGUgZ2lybGllcyBvbiBzdGFuZGJ5IHdhdmluZyBq
dXN0IHRvIHNheSBoaQpEaWQgeW91IHN0b3A/IE5vLCBJIGp1c3QgZHJvdmUg
YnkK".to_vec();
            let secret = base64::decode_config(&secret, base64::MIME).expect("Base64 malformated");
            MyEncrypter {
                suffix: secret,
                key: b"LorentzTransform".to_vec()
            }
        }

        pub fn encrypt(&mut self, buf: &[u8]) -> Vec<u8> {
            let mut plaintext = vec![];
            plaintext.extend_from_slice(buf);
            plaintext.extend_from_slice(&self.suffix);

            //println!("encrypting {:?}", plaintext);
            let cipher = Cipher::aes_128_ecb();
            return symm::encrypt(
                cipher,
                &self.key,
                None,
                &plaintext).unwrap();
        }
    }

    #[test]
    fn test_ecb_decypt_appended() {
        let mut encrypter = MyEncrypter::new();
        let res = ecb_decypt_appended(|buf| encrypter.encrypt(buf));
        // The answer starts with "Rollin' in my", not showing the full one here as a spoiler
        assert!(res.starts_with(&b"Rollin' in my"[..]));
    }

    #[test]
    fn test_pkcs_7_not_multiple() {
        let mut block = b"Minkowski".to_vec();
        pkcs_7(&mut block, 16);
        assert_eq!(&block, b"Minkowski\x07\x07\x07\x07\x07\x07\x07");
    }

    #[test]
    fn test_pkcs_7_multiple() {
        let mut block = b"Euclid".to_vec();
        pkcs_7(&mut block, 6);
        assert_eq!(&block, b"Euclid\x06\x06\x06\x06\x06\x06");
    }
}
<|endoftext|>