<commit_before>#![cfg(test)]
use std::collections::HashMap;

use traits::Cell;
use traits::Grid;
use grid::nhood::MooreNhood;
use grid::square::SquareGrid;

#[derive(Clone, Debug)]
struct MooreTestCell;

impl Cell for MooreTestCell {

    fn step<'a, I>(&self, neighbors: I) -> Self
        where I: Iterator<Item=Option<&'a Self>> {

        // Initial value should be copied
        // There should be 3 neighbors and 5 None values
        // Overall count of neighbors should be equal to 8
        let mut total = 0;
        let mut none_cnt = 0;
        let mut neighbors_cnt = 0;
        for neighbor in neighbors {
            match neighbor {
                None => {
                    none_cnt += 1;
                },
                Some(_) => {
                    neighbors_cnt += 1;
                }
            };
            total += 1;
        }

        assert_eq!(total, 8);
        assert_eq!(none_cnt, 5);
        assert_eq!(neighbors_cnt, 3);

        self.clone()
    }

    fn repr(&self, _: &mut HashMap<&str, &str>) {}
    fn from_repr(&mut self, _: &HashMap<&str, &str>) {}
}

impl Default for MooreTestCell {

    fn default() -> Self { MooreTestCell }
}

#[test]
fn test_neighbors() {
    let nhood = MooreNhood::new();
    let mut grid: SquareGrid<MooreTestCell, _> = SquareGrid::new(2, 2, nhood);
    grid.step();
}
<commit_msg>Refactored to fit well with new repr system.<commit_after>#![cfg(test)]
use std::collections::HashMap;

use traits::Cell;
use traits::Coord;
use traits::Grid;
use traits::Binary;
use grid::nhood::MooreNhood;
use grid::square::SquareGrid;

#[derive(Clone, Debug)]
struct MooreTestCell {
    coord: (i32, i32)
}

impl Cell for MooreTestCell {
    type Coord = (i32, i32);

    fn step<'a, I>(&self, neighbors: I) -> Self
        where I: Iterator<Item=Option<&'a Self>> {

        // Initial value should be copied
        // There should be 3 neighbors and 5 None values
        // Overall count of neighbors should be equal to 8
        let mut total = 0;
        let mut none_cnt = 0;
        let mut neighbors_cnt = 0;
        for neighbor in neighbors {
            match neighbor {
                None => {
                    none_cnt += 1;
                },
                Some(_) => {
                    neighbors_cnt += 1;
                }
            };
            total += 1;
        }

        assert_eq!(total, 8);
        assert_eq!(none_cnt, 5);
        assert_eq!(neighbors_cnt, 3);

        self.clone()
    }

    fn with_coord<C: Coord>(coord: C) -> Self {
        MooreTestCell { coord: (coord.x(), coord.y()) }
    }

    fn coord(&self) -> &Self::Coord {
        &self.coord
    }
}

impl Binary for MooreTestCell {

    fn binary(_: &[u8]) -> Self {
        MooreTestCell { coord: (0, 0) }
    }

    fn bytes(&self) -> &[u8] {

    }
}

#[test]
fn test_neighbors() {
    let nhood = MooreNhood::new();
    let mut grid: SquareGrid<MooreTestCell, _> = SquareGrid::new(2, 2, nhood);
    grid.step();
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate ron;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate euclid;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate error_chain;
extern crate calx_alg;
extern crate calx_grid;
#[macro_use]
extern crate calx_ecs;

use calx_alg::Deciban;

/// Helper macro for formatting textual event messages.
macro_rules! msg {
    ($ctx: expr, $fmt:expr) => {
        let __event = Event::Msg($fmt.to_string());
        $ctx.push_event(__event);
    };

    ($ctx: expr, $fmt:expr, $($arg:expr),*) => {
        let __event = Event::Msg(format!($fmt, $($arg),*));
        $ctx.push_event(__event);
    };
}

mod command;
pub use command::{Command, CommandResult};

mod components;
pub use components::Icon;

mod effect;

mod event;
pub use event::Event;

mod field;
mod flags;

mod form;
pub use form::{Form, FORMS};

mod fov;

mod item;
pub use item::{Slot, ItemType};

mod location;
pub use location::{Location, Portal, Sector};

mod location_set;

mod mapfile;
pub use mapfile::{save_prefab, load_prefab};

mod mutate;
pub use mutate::Mutate;

mod query;
pub use query::Query;

mod spatial;
mod stats;

mod terraform;
pub use terraform::{Terraform, TerrainQuery};

pub mod terrain;
pub use terrain::Terrain;

mod volume;

mod world;
pub use world::{Ecs, World};

mod worldgen;

/// Standard Prefab type, terrain type and spawn name list.
pub type Prefab = calx_grid::Prefab<(Terrain, Vec<String>)>;

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum FovStatus {
    Seen,
    Remembered,
}

pub type Rng = calx_alg::EncodeRng<rand::XorShiftRng>;

/// The combat formula.
///
/// Given a deciban roll and the relevant stats, determine amount of damage dealt.
/// Advantage is attacker skill - target defense.
pub fn attack_damage(roll: f32, advantage: i32, weapon_power: i32) -> i32 {
    const MAX_DAMAGE_MULTIPLIER: f32 = 4.0;

    let roll = roll + advantage as f32;
    (weapon_power as f32 * calx_alg::clamp(0.0, MAX_DAMAGE_MULTIPLIER, (roll - 2.0) * 0.05)) as i32
}

/// Standard deciban roll, clamp into [-20, 20].
pub fn roll<R: rand::Rng>(rng: &mut R) -> f32 {
    calx_alg::clamp(-20.0, 20.0, rng.gen::<Deciban>().0)
}

pub mod errors {
    error_chain! {
        foreign_links {
            Io(::std::io::Error) #[cfg(unix)];
            RonSerialize(::ron::ser::Error);
            RonDeserialize(::ron::de::Error);
        }
    }
}

/// Wrapper class for things that should not be serialized.
struct Cache<T> {
    inner: T,
}

impl<T: Default> Cache<T> {
    pub fn new() -> Cache<T> { Cache { inner: Default::default() } }
}

impl<T> ::std::ops::Deref for Cache<T> {
    type Target = T;

    fn deref(&self) -> &T { &self.inner }
}

impl<T> ::std::ops::DerefMut for Cache<T> {
    fn deref_mut(&mut self) -> &mut T { &mut self.inner }
}
impl<T: Default> serde::Serialize for Cache<T> {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> { ().serialize(s) }
}

impl<'a, T: Default> serde::Deserialize<'a> for Cache<T> {
    fn deserialize<D: serde::Deserializer<'a>>(d: D) -> Result<Self, D::Error> {
        let _: () = serde::Deserialize::deserialize(d)?;
        Ok(Cache::new())
    }
}
<commit_msg>Remove #[cfg(unix)] for std::io::Error foreign_link<commit_after>extern crate rand;
extern crate ron;
extern crate serde;
#[macro_use]
extern crate serde_derive;
extern crate euclid;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate error_chain;
extern crate calx_alg;
extern crate calx_grid;
#[macro_use]
extern crate calx_ecs;

use calx_alg::Deciban;

/// Helper macro for formatting textual event messages.
macro_rules! msg {
    ($ctx: expr, $fmt:expr) => {
        let __event = Event::Msg($fmt.to_string());
        $ctx.push_event(__event);
    };

    ($ctx: expr, $fmt:expr, $($arg:expr),*) => {
        let __event = Event::Msg(format!($fmt, $($arg),*));
        $ctx.push_event(__event);
    };
}

mod command;
pub use command::{Command, CommandResult};

mod components;
pub use components::Icon;

mod effect;

mod event;
pub use event::Event;

mod field;
mod flags;

mod form;
pub use form::{Form, FORMS};

mod fov;

mod item;
pub use item::{Slot, ItemType};

mod location;
pub use location::{Location, Portal, Sector};

mod location_set;

mod mapfile;
pub use mapfile::{save_prefab, load_prefab};

mod mutate;
pub use mutate::Mutate;

mod query;
pub use query::Query;

mod spatial;
mod stats;

mod terraform;
pub use terraform::{Terraform, TerrainQuery};

pub mod terrain;
pub use terrain::Terrain;

mod volume;

mod world;
pub use world::{Ecs, World};

mod worldgen;

/// Standard Prefab type, terrain type and spawn name list.
pub type Prefab = calx_grid::Prefab<(Terrain, Vec<String>)>;

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum FovStatus {
    Seen,
    Remembered,
}

pub type Rng = calx_alg::EncodeRng<rand::XorShiftRng>;

/// The combat formula.
///
/// Given a deciban roll and the relevant stats, determine amount of damage dealt.
/// Advantage is attacker skill - target defense.
pub fn attack_damage(roll: f32, advantage: i32, weapon_power: i32) -> i32 {
    const MAX_DAMAGE_MULTIPLIER: f32 = 4.0;

    let roll = roll + advantage as f32;
    (weapon_power as f32 * calx_alg::clamp(0.0, MAX_DAMAGE_MULTIPLIER, (roll - 2.0) * 0.05)) as i32
}

/// Standard deciban roll, clamp into [-20, 20].
pub fn roll<R: rand::Rng>(rng: &mut R) -> f32 {
    calx_alg::clamp(-20.0, 20.0, rng.gen::<Deciban>().0)
}

pub mod errors {
    error_chain! {
        foreign_links {
            Io(::std::io::Error);
            RonSerialize(::ron::ser::Error);
            RonDeserialize(::ron::de::Error);
        }
    }
}

/// Wrapper class for things that should not be serialized.
struct Cache<T> {
    inner: T,
}

impl<T: Default> Cache<T> {
    pub fn new() -> Cache<T> { Cache { inner: Default::default() } }
}

impl<T> ::std::ops::Deref for Cache<T> {
    type Target = T;

    fn deref(&self) -> &T { &self.inner }
}

impl<T> ::std::ops::DerefMut for Cache<T> {
    fn deref_mut(&mut self) -> &mut T { &mut self.inner }
}
impl<T: Default> serde::Serialize for Cache<T> {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> { ().serialize(s) }
}

impl<'a, T: Default> serde::Deserialize<'a> for Cache<T> {
    fn deserialize<D: serde::Deserializer<'a>>(d: D) -> Result<Self, D::Error> {
        let _: () = serde::Deserialize::deserialize(d)?;
        Ok(Cache::new())
    }
}
<|endoftext|>


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

//! The purpose of KeyManager is to rotate consensus key (and eventually the network key). It is not
//! responsible for generating the first key and fails if the stores have not been properly setup.
//! During rotation, it first updates the local store, then submits a transaction to rotate to the
//! new key. After some period of time and upon restarts of the process, it will evaluate the
//! current status of the system including:
//! * last rotation time, and rotate if it is too long ago
//! * if the latest key in the store matches the latest key in the ValidatorConfig, upon mismatch
//! it will try to submit a transaction to update the ValidatorConfig to the current key in the
//! store.
//! * if the current key in the ValidatorConfig matches the ValidatorSet, if it does not it
//! evaluates the current time from the last reconfiguration and logs that delta with greater
//! levels of severity depending on the delta.
//!
//! KeyManager talks to Libra via the LibraInterface that may either be a direct link into
//! `LibraDB`/`Executor`, JSON-RPC, or some other concoction.
//! KeyManager talks to its own storage through the `LibraSecureStorage::Storage trait.
#![forbid(unsafe_code)]

use crate::{
    libra_interface::LibraInterface,
    logging::{LogEntry, LogEvent, LogSchema},
};
use libra_crypto::ed25519::Ed25519PublicKey;
use libra_global_constants::{CONSENSUS_KEY, OPERATOR_ACCOUNT, OPERATOR_KEY, OWNER_ACCOUNT};
use libra_logger::prelude::*;
use libra_secure_storage::{CryptoStorage, KVStorage};
use libra_secure_time::TimeService;
use libra_types::{
    account_address::AccountAddress,
    account_config::LBR_NAME,
    chain_id::ChainId,
    transaction::{RawTransaction, SignedTransaction, Transaction},
};
use thiserror::Error;

pub mod counters;
pub mod libra_interface;
pub mod logging;

#[cfg(test)]
mod tests;

const GAS_UNIT_PRICE: u64 = 0;
const MAX_GAS_AMOUNT: u64 = 400_000;

/// Defines actions that KeyManager should perform after a check of all associated state.
#[derive(Debug, PartialEq)]
pub enum Action {
    /// The system is in a healthy state and there is no need to perform a rotation
    NoAction,
    /// The system is in a healthy state but sufficient time has passed for another key rotation
    FullKeyRotation,
    /// Storage and the blockchain are inconsistent, submit a new rotation
    SubmitKeyRotationTransaction,
}

#[allow(clippy::large_enum_variant)]
#[derive(Debug, Error, PartialEq)]
pub enum Error {
    #[error("Key mismatch, config: {0}, info: {1}")]
    ConfigInfoKeyMismatch(Ed25519PublicKey, Ed25519PublicKey),
    #[error("Key mismatch, config: {0}, storage: {1}")]
    ConfigStorageKeyMismatch(Ed25519PublicKey, Ed25519PublicKey),
    #[error("Data does not exist: {0}")]
    DataDoesNotExist(String),
    #[error(
        "The libra_timestamp value on-chain isn't increasing. Last value: {0}, Current value: {1}"
    )]
    LivenessError(u64, u64),
    #[error("Unable to retrieve the account address: {0}, storage error: {1}")]
    MissingAccountAddress(String, String),
    #[error("Storage error: {0}")]
    StorageError(String),
    #[error("ValidatorInfo not found in ValidatorConfig: {0}")]
    ValidatorInfoNotFound(AccountAddress),
    #[error("Unknown error: {0}")]
    UnknownError(String),
}

#[cfg(test)]
impl From<anyhow::Error> for Error {
    fn from(error: anyhow::Error) -> Self {
        Error::UnknownError(format!("{}", error))
    }
}

impl From<libra_secure_storage::Error> for Error {
    fn from(error: libra_secure_storage::Error) -> Self {
        Error::StorageError(error.to_string())
    }
}

pub struct KeyManager<LI, S, T> {
    libra: LI,
    storage: S,
    time_service: T,
    last_checked_libra_timestamp: u64,
    rotation_period_secs: u64, // The frequency by which to rotate all keys
    sleep_period_secs: u64,    // The amount of time to sleep between key management checks
    txn_expiration_secs: u64,  // The time after which a rotation transaction expires
    chain_id: ChainId,
}

impl<LI, S, T> KeyManager<LI, S, T>
where
    LI: LibraInterface,
    S: KVStorage + CryptoStorage,
    T: TimeService,
{
    pub fn new(
        libra: LI,
        storage: S,
        time_service: T,
        rotation_period_secs: u64,
        sleep_period_secs: u64,
        txn_expiration_secs: u64,
        chain_id: ChainId,
    ) -> Self {
        Self {
            libra,
            storage,
            time_service,
            last_checked_libra_timestamp: 0,
            rotation_period_secs,
            sleep_period_secs,
            txn_expiration_secs,
            chain_id,
        }
    }

    /// Begins execution of the key manager by running an infinite loop where the key manager will
    /// periodically wake up, verify the state of the validator keys (e.g., the consensus key), and
    /// initiate a key rotation when required. If something goes wrong that we can't handle, an
    /// error will be returned by this method, upon which the key manager will flag the error and
    /// stop execution.
    pub fn execute(&mut self) -> Result<(), Error> {
        loop {
            info!(LogSchema::new(LogEntry::CheckKeyStatus).event(LogEvent::Pending));

            match self.execute_once() {
                Ok(_) => {
                    info!(LogSchema::new(LogEntry::CheckKeyStatus).event(LogEvent::Success));
                }
                Err(Error::LivenessError(last_value, current_value)) => {
                    // Log the liveness error and continue to execute.
                    let error = Error::LivenessError(last_value, current_value);
                    error!(LogSchema::new(LogEntry::CheckKeyStatus)
                        .event(LogEvent::Error)
                        .liveness_error(&error));
                }
                Err(e) => {
                    // Log the unexpected error and continue to execute.
                    error!(LogSchema::new(LogEntry::CheckKeyStatus)
                        .event(LogEvent::Error)
                        .unexpected_error(&e));
                }
            };

            self.sleep();
        }
    }

    fn sleep(&self) {
        info!(LogSchema::new(LogEntry::Sleep)
            .event(LogEvent::Pending)
            .sleep_duration(self.sleep_period_secs));
        self.time_service.sleep(self.sleep_period_secs);

        info!(LogSchema::new(LogEntry::Sleep).event(LogEvent::Success));
    }

    /// Checks the current state of the validator keys and performs any actions that might be
    /// required (e.g., performing a key rotation).
    pub fn execute_once(&mut self) -> Result<(), Error> {
        let action = self.evaluate_status()?;
        self.perform_action(action)
    }

    pub fn compare_storage_to_config(&self) -> Result<(), Error> {
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;

        let storage_key = self.storage.get_public_key(CONSENSUS_KEY)?.public_key;
        let config_key = validator_config.consensus_public_key;
        if storage_key != config_key {
            return Err(Error::ConfigStorageKeyMismatch(config_key, storage_key));
        }

        Ok(())
    }

    pub fn compare_info_to_config(&self) -> Result<(), Error> {
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;
        let validator_info = self.libra.retrieve_validator_info(owner_account)?;

        let info_key = validator_info.consensus_public_key();
        let config_key = validator_config.consensus_public_key;
        if &config_key != info_key {
            return Err(Error::ConfigInfoKeyMismatch(config_key, info_key.clone()));
        }

        Ok(())
    }

    pub fn last_reconfiguration(&self) -> Result<u64, Error> {
        // Convert the time to seconds
        Ok(self.libra.last_reconfiguration()? / 1_000_000)
    }

    pub fn last_rotation(&self) -> Result<u64, Error> {
        Ok(self.storage.get_public_key(CONSENSUS_KEY)?.last_update)
    }

    pub fn libra_timestamp(&self) -> Result<u64, Error> {
        // Convert the time to seconds
        Ok(self.libra.libra_timestamp()? / 1_000_000)
    }

    pub fn resubmit_consensus_key_transaction(&mut self) -> Result<(), Error> {
        let consensus_key = self.storage.get_public_key(CONSENSUS_KEY)?.public_key;
        counters::increment_state("consensus_key", "submitted_rotation_transaction");
        self.submit_key_rotation_transaction(consensus_key)
            .map(|_| ())
    }

    pub fn rotate_consensus_key(&mut self) -> Result<Ed25519PublicKey, Error> {
        let consensus_key = self.storage.rotate_key(CONSENSUS_KEY)?;
        info!(LogSchema::new(LogEntry::KeyRotatedInStorage)
            .event(LogEvent::Success)
            .consensus_key(&consensus_key));
        counters::increment_state("consensus_key", "rotated_in_storage");
        self.submit_key_rotation_transaction(consensus_key)
    }

    pub fn submit_key_rotation_transaction(
        &mut self,
        consensus_key: Ed25519PublicKey,
    ) -> Result<Ed25519PublicKey, Error> {
        let operator_account = self.get_account_from_storage(OPERATOR_ACCOUNT)?;
        let seq_id = self.libra.retrieve_sequence_number(operator_account)?;
        let expiration = self.time_service.now() + self.txn_expiration_secs;

        // Retrieve existing network information as registered on-chain
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;

        let txn = build_rotation_transaction(
            owner_account,
            operator_account,
            seq_id,
            &consensus_key,
            validator_config.validator_network_addresses,
            validator_config.fullnode_network_addresses,
            expiration,
            self.chain_id,
        );

        let operator_pubkey = self.storage.get_public_key(OPERATOR_KEY)?.public_key;
        let txn_signature = self.storage.sign(OPERATOR_KEY, &txn)?;
        let signed_txn = SignedTransaction::new(txn, operator_pubkey, txn_signature);

        self.libra
            .submit_transaction(Transaction::UserTransaction(signed_txn))?;
        info!(LogSchema::new(LogEntry::TransactionSubmission).event(LogEvent::Success));

        Ok(consensus_key)
    }

    /// Ensures that the libra_timestamp() value registered on-chain is strictly monotonically
    /// increasing.
    fn ensure_timestamp_progress(&mut self) -> Result<(), Error> {
        let current_libra_timestamp = self.libra.libra_timestamp()?;
        if current_libra_timestamp <= self.last_checked_libra_timestamp {
            return Err(Error::LivenessError(
                self.last_checked_libra_timestamp,
                current_libra_timestamp,
            ));
        }

        self.last_checked_libra_timestamp = current_libra_timestamp;
        Ok(())
    }

    /// Evaluates the current status of the key manager by performing various state checks between
    /// secure storage and the blockchain.
    ///
    /// Note: every time this function is called, the libra_timestamp registered on-chain must be
    /// strictly monotonically increasing. This helps to ensure that the blockchain is making
    /// progress. Otherwise, if no progress is being made on-chain, a reconfiguration event is
    /// unlikely, and the key manager will be unable to rotate keys.
    pub fn evaluate_status(&mut self) -> Result<Action, Error> {
        self.ensure_timestamp_progress()?;

        // If this is inconsistent, then we are waiting on a reconfiguration...
        if let Err(Error::ConfigInfoKeyMismatch(..)) = self.compare_info_to_config() {
            warn!(LogSchema::new(LogEntry::WaitForReconfiguration));
            counters::increment_state("check_keys", "waiting_on_reconfiguration");
            return Ok(Action::NoAction);
        }

        let last_rotation = self.last_rotation()?;

        // If this is inconsistent, then the transaction either failed or was never submitted.
        if let Err(Error::ConfigStorageKeyMismatch(..)) = self.compare_storage_to_config() {
            return if last_rotation + self.txn_expiration_secs <= self.time_service.now() {
                Ok(Action::SubmitKeyRotationTransaction)
            } else {
                warn!(LogSchema::new(LogEntry::WaitForTransactionExecution));
                counters::increment_state("check_keys", "waiting_on_transaction_execution");
                Ok(Action::NoAction)
            };
        }

        if last_rotation + self.rotation_period_secs <= self.time_service.now() {
            Ok(Action::FullKeyRotation)
        } else {
            info!(LogSchema::new(LogEntry::KeyStillFresh));
            counters::increment_state("check_keys", "keys_still_fresh");
            Ok(Action::NoAction)
        }
    }

    pub fn perform_action(&mut self, action: Action) -> Result<(), Error> {
        match action {
            Action::FullKeyRotation => {
                info!(LogSchema::new(LogEntry::FullKeyRotation).event(LogEvent::Pending));
                self.rotate_consensus_key().map(|_| ())
            }
            Action::SubmitKeyRotationTransaction => {
                info!(LogSchema::new(LogEntry::TransactionSubmission).event(LogEvent::Pending));
                self.resubmit_consensus_key_transaction()
            }
            Action::NoAction => {
                info!(LogSchema::new(LogEntry::NoAction));
                counters::increment_state("check_keys", "no_action");
                Ok(())
            }
        }
    }

    fn get_account_from_storage(&self, account_name: &str) -> Result<AccountAddress, Error> {
        self.storage
            .get::<AccountAddress>(account_name)
            .map(|v| v.value)
            .map_err(|e| Error::MissingAccountAddress(account_name.into(), e.to_string()))
    }
}

pub fn build_rotation_transaction(
    owner_address: AccountAddress,
    operator_address: AccountAddress,
    seq_id: u64,
    consensus_key: &Ed25519PublicKey,
    network_addresses: Vec<u8>,
    fullnode_network_addresses: Vec<u8>,
    expiration_timestamp_secs: u64,
    chain_id: ChainId,
) -> RawTransaction {
    let script =
        transaction_builder_generated::stdlib::encode_set_validator_config_and_reconfigure_script(
            owner_address,
            consensus_key.to_bytes().to_vec(),
            network_addresses,
            fullnode_network_addresses,
        );
    RawTransaction::new_script(
        operator_address,
        seq_id,
        script,
        MAX_GAS_AMOUNT,
        GAS_UNIT_PRICE,
        LBR_NAME.to_owned(),
        expiration_timestamp_secs,
        chain_id,
    )
}
<commit_msg>[Key Manager] Add metrics for the errors encountered by the key manager.<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! The purpose of KeyManager is to rotate consensus key (and eventually the network key). It is not
//! responsible for generating the first key and fails if the stores have not been properly setup.
//! During rotation, it first updates the local store, then submits a transaction to rotate to the
//! new key. After some period of time and upon restarts of the process, it will evaluate the
//! current status of the system including:
//! * last rotation time, and rotate if it is too long ago
//! * if the latest key in the store matches the latest key in the ValidatorConfig, upon mismatch
//! it will try to submit a transaction to update the ValidatorConfig to the current key in the
//! store.
//! * if the current key in the ValidatorConfig matches the ValidatorSet, if it does not it
//! evaluates the current time from the last reconfiguration and logs that delta with greater
//! levels of severity depending on the delta.
//!
//! KeyManager talks to Libra via the LibraInterface that may either be a direct link into
//! `LibraDB`/`Executor`, JSON-RPC, or some other concoction.
//! KeyManager talks to its own storage through the `LibraSecureStorage::Storage trait.
#![forbid(unsafe_code)]

use crate::{
    libra_interface::LibraInterface,
    logging::{LogEntry, LogEvent, LogSchema},
};
use libra_crypto::ed25519::Ed25519PublicKey;
use libra_global_constants::{CONSENSUS_KEY, OPERATOR_ACCOUNT, OPERATOR_KEY, OWNER_ACCOUNT};
use libra_logger::prelude::*;
use libra_secure_storage::{CryptoStorage, KVStorage};
use libra_secure_time::TimeService;
use libra_types::{
    account_address::AccountAddress,
    account_config::LBR_NAME,
    chain_id::ChainId,
    transaction::{RawTransaction, SignedTransaction, Transaction},
};
use thiserror::Error;

pub mod counters;
pub mod libra_interface;
pub mod logging;

#[cfg(test)]
mod tests;

const GAS_UNIT_PRICE: u64 = 0;
const MAX_GAS_AMOUNT: u64 = 400_000;

/// Defines actions that KeyManager should perform after a check of all associated state.
#[derive(Debug, PartialEq)]
pub enum Action {
    /// The system is in a healthy state and there is no need to perform a rotation
    NoAction,
    /// The system is in a healthy state but sufficient time has passed for another key rotation
    FullKeyRotation,
    /// Storage and the blockchain are inconsistent, submit a new rotation
    SubmitKeyRotationTransaction,
}

#[allow(clippy::large_enum_variant)]
#[derive(Debug, Error, PartialEq)]
pub enum Error {
    #[error("Key mismatch, config: {0}, info: {1}")]
    ConfigInfoKeyMismatch(Ed25519PublicKey, Ed25519PublicKey),
    #[error("Key mismatch, config: {0}, storage: {1}")]
    ConfigStorageKeyMismatch(Ed25519PublicKey, Ed25519PublicKey),
    #[error("Data does not exist: {0}")]
    DataDoesNotExist(String),
    #[error(
        "The libra_timestamp value on-chain isn't increasing. Last value: {0}, Current value: {1}"
    )]
    LivenessError(u64, u64),
    #[error("Unable to retrieve the account address: {0}, storage error: {1}")]
    MissingAccountAddress(String, String),
    #[error("Storage error: {0}")]
    StorageError(String),
    #[error("ValidatorInfo not found in ValidatorConfig: {0}")]
    ValidatorInfoNotFound(AccountAddress),
    #[error("Unknown error: {0}")]
    UnknownError(String),
}

#[cfg(test)]
impl From<anyhow::Error> for Error {
    fn from(error: anyhow::Error) -> Self {
        Error::UnknownError(format!("{}", error))
    }
}

impl From<libra_secure_storage::Error> for Error {
    fn from(error: libra_secure_storage::Error) -> Self {
        Error::StorageError(error.to_string())
    }
}

pub struct KeyManager<LI, S, T> {
    libra: LI,
    storage: S,
    time_service: T,
    last_checked_libra_timestamp: u64,
    rotation_period_secs: u64, // The frequency by which to rotate all keys
    sleep_period_secs: u64,    // The amount of time to sleep between key management checks
    txn_expiration_secs: u64,  // The time after which a rotation transaction expires
    chain_id: ChainId,
}

impl<LI, S, T> KeyManager<LI, S, T>
where
    LI: LibraInterface,
    S: KVStorage + CryptoStorage,
    T: TimeService,
{
    pub fn new(
        libra: LI,
        storage: S,
        time_service: T,
        rotation_period_secs: u64,
        sleep_period_secs: u64,
        txn_expiration_secs: u64,
        chain_id: ChainId,
    ) -> Self {
        Self {
            libra,
            storage,
            time_service,
            last_checked_libra_timestamp: 0,
            rotation_period_secs,
            sleep_period_secs,
            txn_expiration_secs,
            chain_id,
        }
    }

    /// Begins execution of the key manager by running an infinite loop where the key manager will
    /// periodically wake up, verify the state of the validator keys (e.g., the consensus key), and
    /// initiate a key rotation when required. If something goes wrong that we can't handle, an
    /// error will be returned by this method, upon which the key manager will flag the error and
    /// stop execution.
    pub fn execute(&mut self) -> Result<(), Error> {
        loop {
            info!(LogSchema::new(LogEntry::CheckKeyStatus).event(LogEvent::Pending));

            match self.execute_once() {
                Ok(_) => {
                    info!(LogSchema::new(LogEntry::CheckKeyStatus).event(LogEvent::Success));
                }
                Err(Error::LivenessError(last_value, current_value)) => {
                    // Log the liveness error and continue to execute.
                    let error = Error::LivenessError(last_value, current_value);
                    error!(LogSchema::new(LogEntry::CheckKeyStatus)
                        .event(LogEvent::Error)
                        .liveness_error(&error));
                    counters::increment_state("check_keys", "liveness_error_encountered");
                }
                Err(e) => {
                    // Log the unexpected error and continue to execute.
                    error!(LogSchema::new(LogEntry::CheckKeyStatus)
                        .event(LogEvent::Error)
                        .unexpected_error(&e));
                    counters::increment_state("check_keys", "unexpected_error_encountered");
                }
            };

            self.sleep();
        }
    }

    fn sleep(&self) {
        info!(LogSchema::new(LogEntry::Sleep)
            .event(LogEvent::Pending)
            .sleep_duration(self.sleep_period_secs));
        self.time_service.sleep(self.sleep_period_secs);

        info!(LogSchema::new(LogEntry::Sleep).event(LogEvent::Success));
    }

    /// Checks the current state of the validator keys and performs any actions that might be
    /// required (e.g., performing a key rotation).
    pub fn execute_once(&mut self) -> Result<(), Error> {
        let action = self.evaluate_status()?;
        self.perform_action(action)
    }

    pub fn compare_storage_to_config(&self) -> Result<(), Error> {
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;

        let storage_key = self.storage.get_public_key(CONSENSUS_KEY)?.public_key;
        let config_key = validator_config.consensus_public_key;
        if storage_key != config_key {
            return Err(Error::ConfigStorageKeyMismatch(config_key, storage_key));
        }

        Ok(())
    }

    pub fn compare_info_to_config(&self) -> Result<(), Error> {
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;
        let validator_info = self.libra.retrieve_validator_info(owner_account)?;

        let info_key = validator_info.consensus_public_key();
        let config_key = validator_config.consensus_public_key;
        if &config_key != info_key {
            return Err(Error::ConfigInfoKeyMismatch(config_key, info_key.clone()));
        }

        Ok(())
    }

    pub fn last_reconfiguration(&self) -> Result<u64, Error> {
        // Convert the time to seconds
        Ok(self.libra.last_reconfiguration()? / 1_000_000)
    }

    pub fn last_rotation(&self) -> Result<u64, Error> {
        Ok(self.storage.get_public_key(CONSENSUS_KEY)?.last_update)
    }

    pub fn libra_timestamp(&self) -> Result<u64, Error> {
        // Convert the time to seconds
        Ok(self.libra.libra_timestamp()? / 1_000_000)
    }

    pub fn resubmit_consensus_key_transaction(&mut self) -> Result<(), Error> {
        let consensus_key = self.storage.get_public_key(CONSENSUS_KEY)?.public_key;
        counters::increment_state("consensus_key", "submitted_rotation_transaction");
        self.submit_key_rotation_transaction(consensus_key)
            .map(|_| ())
    }

    pub fn rotate_consensus_key(&mut self) -> Result<Ed25519PublicKey, Error> {
        let consensus_key = self.storage.rotate_key(CONSENSUS_KEY)?;
        info!(LogSchema::new(LogEntry::KeyRotatedInStorage)
            .event(LogEvent::Success)
            .consensus_key(&consensus_key));
        counters::increment_state("consensus_key", "rotated_in_storage");
        self.submit_key_rotation_transaction(consensus_key)
    }

    pub fn submit_key_rotation_transaction(
        &mut self,
        consensus_key: Ed25519PublicKey,
    ) -> Result<Ed25519PublicKey, Error> {
        let operator_account = self.get_account_from_storage(OPERATOR_ACCOUNT)?;
        let seq_id = self.libra.retrieve_sequence_number(operator_account)?;
        let expiration = self.time_service.now() + self.txn_expiration_secs;

        // Retrieve existing network information as registered on-chain
        let owner_account = self.get_account_from_storage(OWNER_ACCOUNT)?;
        let validator_config = self.libra.retrieve_validator_config(owner_account)?;

        let txn = build_rotation_transaction(
            owner_account,
            operator_account,
            seq_id,
            &consensus_key,
            validator_config.validator_network_addresses,
            validator_config.fullnode_network_addresses,
            expiration,
            self.chain_id,
        );

        let operator_pubkey = self.storage.get_public_key(OPERATOR_KEY)?.public_key;
        let txn_signature = self.storage.sign(OPERATOR_KEY, &txn)?;
        let signed_txn = SignedTransaction::new(txn, operator_pubkey, txn_signature);

        self.libra
            .submit_transaction(Transaction::UserTransaction(signed_txn))?;
        info!(LogSchema::new(LogEntry::TransactionSubmission).event(LogEvent::Success));

        Ok(consensus_key)
    }

    /// Ensures that the libra_timestamp() value registered on-chain is strictly monotonically
    /// increasing.
    fn ensure_timestamp_progress(&mut self) -> Result<(), Error> {
        let current_libra_timestamp = self.libra.libra_timestamp()?;
        if current_libra_timestamp <= self.last_checked_libra_timestamp {
            return Err(Error::LivenessError(
                self.last_checked_libra_timestamp,
                current_libra_timestamp,
            ));
        }

        self.last_checked_libra_timestamp = current_libra_timestamp;
        Ok(())
    }

    /// Evaluates the current status of the key manager by performing various state checks between
    /// secure storage and the blockchain.
    ///
    /// Note: every time this function is called, the libra_timestamp registered on-chain must be
    /// strictly monotonically increasing. This helps to ensure that the blockchain is making
    /// progress. Otherwise, if no progress is being made on-chain, a reconfiguration event is
    /// unlikely, and the key manager will be unable to rotate keys.
    pub fn evaluate_status(&mut self) -> Result<Action, Error> {
        self.ensure_timestamp_progress()?;

        // If this is inconsistent, then we are waiting on a reconfiguration...
        if let Err(Error::ConfigInfoKeyMismatch(..)) = self.compare_info_to_config() {
            warn!(LogSchema::new(LogEntry::WaitForReconfiguration));
            counters::increment_state("check_keys", "waiting_on_reconfiguration");
            return Ok(Action::NoAction);
        }

        let last_rotation = self.last_rotation()?;

        // If this is inconsistent, then the transaction either failed or was never submitted.
        if let Err(Error::ConfigStorageKeyMismatch(..)) = self.compare_storage_to_config() {
            return if last_rotation + self.txn_expiration_secs <= self.time_service.now() {
                Ok(Action::SubmitKeyRotationTransaction)
            } else {
                warn!(LogSchema::new(LogEntry::WaitForTransactionExecution));
                counters::increment_state("check_keys", "waiting_on_transaction_execution");
                Ok(Action::NoAction)
            };
        }

        if last_rotation + self.rotation_period_secs <= self.time_service.now() {
            Ok(Action::FullKeyRotation)
        } else {
            info!(LogSchema::new(LogEntry::KeyStillFresh));
            counters::increment_state("check_keys", "keys_still_fresh");
            Ok(Action::NoAction)
        }
    }

    pub fn perform_action(&mut self, action: Action) -> Result<(), Error> {
        match action {
            Action::FullKeyRotation => {
                info!(LogSchema::new(LogEntry::FullKeyRotation).event(LogEvent::Pending));
                self.rotate_consensus_key().map(|_| ())
            }
            Action::SubmitKeyRotationTransaction => {
                info!(LogSchema::new(LogEntry::TransactionSubmission).event(LogEvent::Pending));
                self.resubmit_consensus_key_transaction()
            }
            Action::NoAction => {
                info!(LogSchema::new(LogEntry::NoAction));
                counters::increment_state("check_keys", "no_action");
                Ok(())
            }
        }
    }

    fn get_account_from_storage(&self, account_name: &str) -> Result<AccountAddress, Error> {
        self.storage
            .get::<AccountAddress>(account_name)
            .map(|v| v.value)
            .map_err(|e| Error::MissingAccountAddress(account_name.into(), e.to_string()))
    }
}

pub fn build_rotation_transaction(
    owner_address: AccountAddress,
    operator_address: AccountAddress,
    seq_id: u64,
    consensus_key: &Ed25519PublicKey,
    network_addresses: Vec<u8>,
    fullnode_network_addresses: Vec<u8>,
    expiration_timestamp_secs: u64,
    chain_id: ChainId,
) -> RawTransaction {
    let script =
        transaction_builder_generated::stdlib::encode_set_validator_config_and_reconfigure_script(
            owner_address,
            consensus_key.to_bytes().to_vec(),
            network_addresses,
            fullnode_network_addresses,
        );
    RawTransaction::new_script(
        operator_address,
        seq_id,
        script,
        MAX_GAS_AMOUNT,
        GAS_UNIT_PRICE,
        LBR_NAME.to_owned(),
        expiration_timestamp_secs,
        chain_id,
    )
}
<|endoftext|>


<commit_before>//! Client handshake machine.

use std::io::{Read, Write};
use std::marker::PhantomData;

use base64;
use httparse::Status;
use httparse;
use rand;
use url::Url;

use error::{Error, Result};
use protocol::{WebSocket, Role};
use super::headers::{Headers, FromHttparse, MAX_HEADERS};
use super::machine::{HandshakeMachine, StageResult, TryParse};
use super::{MidHandshake, HandshakeRole, ProcessingResult, convert_key};

/// Client request.
pub struct Request<'t> {
    /// `ws://` or `wss://` URL to connect to.
    pub url: Url,
    /// Extra HTTP headers to append to the request.
    pub extra_headers: Option<&'t [(&'t str, &'t str)]>,
}

impl<'t> Request<'t> {
    /// The GET part of the request.
    fn get_path(&self) -> String {
        if let Some(query) = self.url.query() {
            format!("{path}?{query}", path = self.url.path(), query = query)
        } else {
            self.url.path().into()
        }
    }
    /// The Host: part of the request.
    fn get_host(&self) -> String {
        let host = self.url.host_str().expect("Bug: URL without host");
        if let Some(port) = self.url.port() {
            format!("{host}:{port}", host = host, port = port)
        } else {
            host.into()
        }
    }
}

impl From<Url> for Request<'static> {
    fn from(value: Url) -> Self {
        Request {
            url: value,
            extra_headers: None,
        }
    }
}

/// Client handshake role.
pub struct ClientHandshake<S> {
    verify_data: VerifyData,
    _marker: PhantomData<S>,
}

impl<S: Read + Write> ClientHandshake<S> {
    /// Initiate a client handshake.
    pub fn start(stream: S, request: Request) -> MidHandshake<Self> {
        let key = generate_key();

        let machine = {
            let mut req = Vec::new();
            write!(req, "\
                GET {path} HTTP/1.1\r\n\
                Host: {host}\r\n\
                Connection: upgrade\r\n\
                Upgrade: websocket\r\n\
                Sec-WebSocket-Version: 13\r\n\
                Sec-WebSocket-Key: {key}\r\n",
                host = request.get_host(), path = request.get_path(), key = key).unwrap();
            if let Some(eh) = request.extra_headers {
                for &(k, v) in eh {
                    write!(req, "{}: {}\r\n", k, v).unwrap();
                }
            }
            write!(req, "\r\n").unwrap();
            HandshakeMachine::start_write(stream, req)
        };

        let client = {
            let accept_key = convert_key(key.as_ref()).unwrap();
            ClientHandshake {
                verify_data: VerifyData { accept_key },
                _marker: PhantomData,
            }
        };

        trace!("Client handshake initiated.");
        MidHandshake { role: client, machine: machine }
    }
}

impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
    type IncomingData = Response;
    type InternalStream = S;
    type FinalResult = (WebSocket<S>, Response);
    fn stage_finished(&mut self, finish: StageResult<Self::IncomingData, Self::InternalStream>)
        -> Result<ProcessingResult<Self::InternalStream, Self::FinalResult>>
    {
        Ok(match finish {
            StageResult::DoneWriting(stream) => {
                ProcessingResult::Continue(HandshakeMachine::start_read(stream))
            }
            StageResult::DoneReading { stream, result, tail, } => {
                self.verify_data.verify_response(&result)?;
                debug!("Client handshake done.");
                ProcessingResult::Done((WebSocket::from_partially_read(stream, tail, Role::Client),
                                       result))
            }
        })
    }
}

/// Information for handshake verification.
struct VerifyData {
    /// Accepted server key.
    accept_key: String,
}

impl VerifyData {
    pub fn verify_response(&self, response: &Response) -> Result<()> {
        // 1. If the status code received from the server is not 101, the
        // client handles the response per HTTP [RFC2616] procedures. (RFC 6455)
        if response.code != 101 {
            return Err(Error::Http(response.code));
        }
        // 2. If the response lacks an |Upgrade| header field or the |Upgrade|
        // header field contains a value that is not an ASCII case-
        // insensitive match for the value "websocket", the client MUST
        // _Fail the WebSocket Connection_. (RFC 6455)
        if !response.headers.header_is_ignore_case("Upgrade", "websocket") {
            return Err(Error::Protocol("No \"Upgrade: websocket\" in server reply".into()));
        }
        // 3.  If the response lacks a |Connection| header field or the
        // |Connection| header field doesn't contain a token that is an
        // ASCII case-insensitive match for the value "Upgrade", the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        if !response.headers.header_is_ignore_case("Connection", "Upgrade") {
            return Err(Error::Protocol("No \"Connection: upgrade\" in server reply".into()));
        }
        // 4.  If the response lacks a |Sec-WebSocket-Accept| header field or
        // the |Sec-WebSocket-Accept| contains a value other than the
        // base64-encoded SHA-1 of ... the client MUST _Fail the WebSocket
        // Connection_. (RFC 6455)
        if !response.headers.header_is("Sec-WebSocket-Accept", &self.accept_key) {
            return Err(Error::Protocol("Key mismatch in Sec-WebSocket-Accept".into()));
        }
        // 5.  If the response includes a |Sec-WebSocket-Extensions| header
        // field and this header field indicates the use of an extension
        // that was not present in the client's handshake (the server has
        // indicated an extension not requested by the client), the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        // TODO

        // 6.  If the response includes a |Sec-WebSocket-Protocol| header field
        // and this header field indicates the use of a subprotocol that was
        // not present in the client's handshake (the server has indicated a
        // subprotocol not requested by the client), the client MUST _Fail
        // the WebSocket Connection_. (RFC 6455)
        // TODO

        Ok(())
    }
}

/// Server response.
pub struct Response {
    /// HTTP response code of the response.
    pub code: u16,
    /// Received headers.
    pub headers: Headers,
}

impl TryParse for Response {
    fn try_parse(buf: &[u8]) -> Result<Option<(usize, Self)>> {
        let mut hbuffer = [httparse::EMPTY_HEADER; MAX_HEADERS];
        let mut req = httparse::Response::new(&mut hbuffer);
        Ok(match req.parse(buf)? {
            Status::Partial => None,
            Status::Complete(size) => Some((size, Response::from_httparse(req)?)),
        })
    }
}

impl<'h, 'b: 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
    fn from_httparse(raw: httparse::Response<'h, 'b>) -> Result<Self> {
        if raw.version.expect("Bug: no HTTP version") < /*1.*/1 {
            return Err(Error::Protocol("HTTP version should be 1.1 or higher".into()));
        }
        Ok(Response {
            code: raw.code.expect("Bug: no HTTP response code"),
            headers: Headers::from_httparse(raw.headers)?,
        })
    }
}

/// Generate a random key for the `Sec-WebSocket-Key` header.
fn generate_key() -> String {
    // a base64-encoded (see Section 4 of [RFC4648]) value that,
    // when decoded, is 16 bytes in length (RFC 6455)
    let r: [u8; 16] = rand::random();
    base64::encode(&r)
}

#[cfg(test)]
mod tests {
    use super::{Response, generate_key};
    use super::super::machine::TryParse;

    #[test]
    fn random_keys() {
        let k1 = generate_key();
        println!("Generated random key 1: {}", k1);
        let k2 = generate_key();
        println!("Generated random key 2: {}", k2);
        assert_ne!(k1, k2);
        assert_eq!(k1.len(), k2.len());
        assert_eq!(k1.len(), 24);
        assert_eq!(k2.len(), 24);
        assert!(k1.ends_with("=="));
        assert!(k2.ends_with("=="));
        assert!(k1[..22].find("=").is_none());
        assert!(k2[..22].find("=").is_none());
    }

    #[test]
    fn response_parsing() {
        const DATA: &'static [u8] = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n";
        let (_, resp) = Response::try_parse(DATA).unwrap().unwrap();
        assert_eq!(resp.code, 200);
        assert_eq!(resp.headers.find_first("Content-Type"), Some(&b"text/html"[..]));
    }
}
<commit_msg>Improve the `handshake::client::Request` structure<commit_after>//! Client handshake machine.

use std::borrow::Cow;
use std::io::{Read, Write};
use std::marker::PhantomData;

use base64;
use httparse::Status;
use httparse;
use rand;
use url::Url;

use error::{Error, Result};
use protocol::{WebSocket, Role};
use super::headers::{Headers, FromHttparse, MAX_HEADERS};
use super::machine::{HandshakeMachine, StageResult, TryParse};
use super::{MidHandshake, HandshakeRole, ProcessingResult, convert_key};

/// Client request.
pub struct Request<'t> {
    /// `ws://` or `wss://` URL to connect to.
    pub url: Url,
    /// Extra HTTP headers to append to the request.
    pub extra_headers: Option<Vec<(Cow<'t, str>, Cow<'t, str>)>>,
}

impl<'t> Request<'t> {
    /// Returns the GET part of the request.
    fn get_path(&self) -> String {
        if let Some(query) = self.url.query() {
            format!("{path}?{query}", path = self.url.path(), query = query)
        } else {
            self.url.path().into()
        }
    }

    /// Returns the host part of the request.
    fn get_host(&self) -> String {
        let host = self.url.host_str().expect("Bug: URL without host");
        if let Some(port) = self.url.port() {
            format!("{host}:{port}", host = host, port = port)
        } else {
            host.into()
        }
    }

    /// Adds a WebSocket protocol to the request.
    pub fn add_protocol(&mut self, protocol: Cow<'t, str>) {
        self.add_header(Cow::from("Sec-WebSocket-Protocol"), protocol);
    }

    /// Adds a custom header to the request.
    pub fn add_header(&mut self, name: Cow<'t, str>, value: Cow<'t, str>) {
        let mut headers = self.extra_headers.take().unwrap_or(vec![]);
        headers.push((name, value));
        self.extra_headers = Some(headers);
    }
}

impl From<Url> for Request<'static> {
    fn from(value: Url) -> Self {
        Request {
            url: value,
            extra_headers: None,
        }
    }
}

/// Client handshake role.
pub struct ClientHandshake<S> {
    verify_data: VerifyData,
    _marker: PhantomData<S>,
}

impl<S: Read + Write> ClientHandshake<S> {
    /// Initiate a client handshake.
    pub fn start(stream: S, request: Request) -> MidHandshake<Self> {
        let key = generate_key();

        let machine = {
            let mut req = Vec::new();
            write!(req, "\
                GET {path} HTTP/1.1\r\n\
                Host: {host}\r\n\
                Connection: upgrade\r\n\
                Upgrade: websocket\r\n\
                Sec-WebSocket-Version: 13\r\n\
                Sec-WebSocket-Key: {key}\r\n",
                host = request.get_host(), path = request.get_path(), key = key).unwrap();
            if let Some(eh) = request.extra_headers {
                for (k, v) in eh {
                    write!(req, "{}: {}\r\n", k, v).unwrap();
                }
            }
            write!(req, "\r\n").unwrap();
            HandshakeMachine::start_write(stream, req)
        };

        let client = {
            let accept_key = convert_key(key.as_ref()).unwrap();
            ClientHandshake {
                verify_data: VerifyData { accept_key },
                _marker: PhantomData,
            }
        };

        trace!("Client handshake initiated.");
        MidHandshake { role: client, machine: machine }
    }
}

impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
    type IncomingData = Response;
    type InternalStream = S;
    type FinalResult = (WebSocket<S>, Response);
    fn stage_finished(&mut self, finish: StageResult<Self::IncomingData, Self::InternalStream>)
        -> Result<ProcessingResult<Self::InternalStream, Self::FinalResult>>
    {
        Ok(match finish {
            StageResult::DoneWriting(stream) => {
                ProcessingResult::Continue(HandshakeMachine::start_read(stream))
            }
            StageResult::DoneReading { stream, result, tail, } => {
                self.verify_data.verify_response(&result)?;
                debug!("Client handshake done.");
                ProcessingResult::Done((WebSocket::from_partially_read(stream, tail, Role::Client),
                                       result))
            }
        })
    }
}

/// Information for handshake verification.
struct VerifyData {
    /// Accepted server key.
    accept_key: String,
}

impl VerifyData {
    pub fn verify_response(&self, response: &Response) -> Result<()> {
        // 1. If the status code received from the server is not 101, the
        // client handles the response per HTTP [RFC2616] procedures. (RFC 6455)
        if response.code != 101 {
            return Err(Error::Http(response.code));
        }
        // 2. If the response lacks an |Upgrade| header field or the |Upgrade|
        // header field contains a value that is not an ASCII case-
        // insensitive match for the value "websocket", the client MUST
        // _Fail the WebSocket Connection_. (RFC 6455)
        if !response.headers.header_is_ignore_case("Upgrade", "websocket") {
            return Err(Error::Protocol("No \"Upgrade: websocket\" in server reply".into()));
        }
        // 3.  If the response lacks a |Connection| header field or the
        // |Connection| header field doesn't contain a token that is an
        // ASCII case-insensitive match for the value "Upgrade", the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        if !response.headers.header_is_ignore_case("Connection", "Upgrade") {
            return Err(Error::Protocol("No \"Connection: upgrade\" in server reply".into()));
        }
        // 4.  If the response lacks a |Sec-WebSocket-Accept| header field or
        // the |Sec-WebSocket-Accept| contains a value other than the
        // base64-encoded SHA-1 of ... the client MUST _Fail the WebSocket
        // Connection_. (RFC 6455)
        if !response.headers.header_is("Sec-WebSocket-Accept", &self.accept_key) {
            return Err(Error::Protocol("Key mismatch in Sec-WebSocket-Accept".into()));
        }
        // 5.  If the response includes a |Sec-WebSocket-Extensions| header
        // field and this header field indicates the use of an extension
        // that was not present in the client's handshake (the server has
        // indicated an extension not requested by the client), the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        // TODO

        // 6.  If the response includes a |Sec-WebSocket-Protocol| header field
        // and this header field indicates the use of a subprotocol that was
        // not present in the client's handshake (the server has indicated a
        // subprotocol not requested by the client), the client MUST _Fail
        // the WebSocket Connection_. (RFC 6455)
        // TODO

        Ok(())
    }
}

/// Server response.
pub struct Response {
    /// HTTP response code of the response.
    pub code: u16,
    /// Received headers.
    pub headers: Headers,
}

impl TryParse for Response {
    fn try_parse(buf: &[u8]) -> Result<Option<(usize, Self)>> {
        let mut hbuffer = [httparse::EMPTY_HEADER; MAX_HEADERS];
        let mut req = httparse::Response::new(&mut hbuffer);
        Ok(match req.parse(buf)? {
            Status::Partial => None,
            Status::Complete(size) => Some((size, Response::from_httparse(req)?)),
        })
    }
}

impl<'h, 'b: 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
    fn from_httparse(raw: httparse::Response<'h, 'b>) -> Result<Self> {
        if raw.version.expect("Bug: no HTTP version") < /*1.*/1 {
            return Err(Error::Protocol("HTTP version should be 1.1 or higher".into()));
        }
        Ok(Response {
            code: raw.code.expect("Bug: no HTTP response code"),
            headers: Headers::from_httparse(raw.headers)?,
        })
    }
}

/// Generate a random key for the `Sec-WebSocket-Key` header.
fn generate_key() -> String {
    // a base64-encoded (see Section 4 of [RFC4648]) value that,
    // when decoded, is 16 bytes in length (RFC 6455)
    let r: [u8; 16] = rand::random();
    base64::encode(&r)
}

#[cfg(test)]
mod tests {
    use super::{Response, generate_key};
    use super::super::machine::TryParse;

    #[test]
    fn random_keys() {
        let k1 = generate_key();
        println!("Generated random key 1: {}", k1);
        let k2 = generate_key();
        println!("Generated random key 2: {}", k2);
        assert_ne!(k1, k2);
        assert_eq!(k1.len(), k2.len());
        assert_eq!(k1.len(), 24);
        assert_eq!(k2.len(), 24);
        assert!(k1.ends_with("=="));
        assert!(k2.ends_with("=="));
        assert!(k1[..22].find("=").is_none());
        assert!(k2[..22].find("=").is_none());
    }

    #[test]
    fn response_parsing() {
        const DATA: &'static [u8] = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n";
        let (_, resp) = Response::try_parse(DATA).unwrap().unwrap();
        assert_eq!(resp.code, 200);
        assert_eq!(resp.headers.find_first("Content-Type"), Some(&b"text/html"[..]));
    }
}
<|endoftext|>


<commit_before>#[macro_use]
pub mod types;

use self::types::*;
use common::combinator::Parser;

pub fn tokenize<T: Iterator<Item = char>>(input: T) -> Result<Vec<Token>> {
    let mut l = Lexer::new(input);

    let mut result = Vec::default();

    loop {
        let _ = skip_whitespaces(&mut l);
        match l.preview() {
            None => break,
            _ => {
                result.push(try!(lex_token(&mut l)));
            }
        }
    }

    Ok(result)
}

fn skip_whitespaces(l: &mut Lexer) -> Result<Vec<char>> {
    l.many(|l| l.predicate(|x| [' ', '\t', '\n', '\r'].contains(x)))
}

fn lex_token(l: &mut Lexer) -> Result<Token> {
    l.choose(&[&lex_operator, &lex_punctuation])
}

macro_rules! parse_map {
    ($pat:expr, $tok:ident) => {
        &|l| {
            let pos = l.current_pos();
            let lit: String = try!(l.string($pat.chars()));
            Ok(token!($tok, pos.0, pos.1, lit))
        }
    }
}

fn lex_operator(l: &mut Lexer) -> Result<Token> {
    l.choose(&[parse_map!("==", Eq),
               parse_map!("=", Assign),
               parse_map!("+", Plus),
               parse_map!("-", Minus),
               parse_map!("*", Multiply),
               parse_map!("/", Divide),
               parse_map!("!=", NotEq),
               parse_map!("!", Not),
               parse_map!(">", GreaterThan),
               parse_map!("<", LessThan)])
}

fn lex_punctuation(l: &mut Lexer) -> Result<Token> {
    l.choose(&[parse_map!(":", Colon),
               parse_map!(";", SemiColon),
               parse_map!(",", Comma),
               parse_map!("(", LParen),
               parse_map!(")", RParen),
               parse_map!("{", LBrace),
               parse_map!("}", RBrace),
               parse_map!("[", LBracket),
               parse_map!("]", RBracket)])
}

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

    fn lex(input: &str) -> Vec<Token> {
        tokenize(String::from(input).chars()).unwrap()
    }

    fn vec_eq<T: PartialEq + Debug>(xs: &Vec<T>, ys: &Vec<T>) {
        if xs.len() != ys.len() {
            panic!("{:?} != {:?}", xs, ys);
        }
        for i in 0..xs.len() {
            if xs[i] != ys[i] {
                panic!("{:?} != {:?}", xs, ys);
            }
        }
    }

    #[test]
    fn special_chars() {
        vec_eq(&lex("=+(){},;"),
               &vec![token!(Assign, 1, 1, '='),
                     token!(Plus, 1, 2, '+'),
                     token!(LParen, 1, 3, '('),
                     token!(RParen, 1, 4, ')'),
                     token!(LBrace, 1, 5, '{'),
                     token!(RBrace, 1, 6, '}'),
                     token!(Comma, 1, 7, ','),
                     token!(SemiColon, 1, 8, ';')]);
    }

    #[test]
    fn complex_code() {
        vec_eq(&lex("
let five = 5;
let ten = 10;
let add = fn(x, y) {
    x + y;
};
let result = add(five, ten);
               "),
               &vec![token!(Let, 2, 1, "let"),
                     token!(Ident, 2, 5, "five"),
                     token!(Assign, 2, 10, "="),
                     token!(IntLiteral, 2, 12, "5"),
                     token!(SemiColon, 2, 13, ";"),
                     token!(Let, 3, 1, "let"),
                     token!(Ident, 3, 5, "ten"),
                     token!(Assign, 3, 9, "="),
                     token!(IntLiteral, 3, 11, "10"),
                     token!(SemiColon, 3, 13, ";"),
                     token!(Let, 4, 1, "let"),
                     token!(Ident, 4, 5, "add"),
                     token!(Assign, 4, 9, "="),
                     token!(Function, 4, 11, "fn"),
                     token!(LParen, 4, 13, "("),
                     token!(Ident, 4, 14, "x"),
                     token!(Comma, 4, 15, ","),
                     token!(Ident, 4, 17, "y"),
                     token!(RParen, 4, 18, ")"),
                     token!(LBrace, 4, 20, "{"),
                     token!(Ident, 5, 5, "x"),
                     token!(Plus, 5, 7, "+"),
                     token!(Ident, 5, 9, "y"),
                     token!(SemiColon, 5, 10, ";"),
                     token!(RBrace, 6, 1, "}"),
                     token!(SemiColon, 6, 2, ";"),
                     token!(Let, 7, 1, "let"),
                     token!(Ident, 7, 5, "result"),
                     token!(Assign, 7, 12, "="),
                     token!(Ident, 7, 14, "add"),
                     token!(LParen, 7, 17, "("),
                     token!(Ident, 7, 18, "five"),
                     token!(Comma, 7, 22, ","),
                     token!(Ident, 7, 24, "ten"),
                     token!(RParen, 7, 27, ")"),
                     token!(SemiColon, 7, 28, ";")]);
    }

    #[test]
    fn complex_code2() {
        vec_eq(&lex("
if (a == 10) {
  return a;
} else if (a != 20) {
  return !a;
} else if (a > 20) {
  return -30 / 40 * 50;
} else if (a < 30) {
  return true;
}
return false;
                    "),
               &vec![token!(If, 2, 1, "if"),
                     token!(LParen, 2, 3, "("),
                     token!(Ident, 2, 4, "a"),
                     token!(Eq, 2, 6, "=="),
                     token!(IntLiteral, 2, 9, "10"),
                     token!(RParen, 2, 11, ")"),
                     token!(LBrace, 2, 13, "{"),
                     token!(Return, 3, 3, "return"),
                     token!(Ident, 3, 10, "a"),
                     token!(SemiColon, 3, 11, ";"),
                     token!(RBrace, 4, 1, "}"),
                     token!(Else, 4, 3, "else"),
                     token!(If, 4, 8, "if"),
                     token!(LParen, 4, 11, "("),
                     token!(Ident, 4, 12, "a"),
                     token!(NotEq, 4, 14, "!="),
                     token!(IntLiteral, 4, 17, "20"),
                     token!(RParen, 4, 19, ")"),
                     token!(LBrace, 4, 21, "{"),
                     token!(Return, 5, 3, "return"),
                     token!(Not, 5, 10, "!"),
                     token!(Ident, 5, 11, "a"),
                     token!(SemiColon, 5, 13, ";"),
                     token!(RBrace, 6, 1, "}"),
                     token!(Else, 6, 3, "else"),
                     token!(If, 6, 8, "if"),
                     token!(LParen, 6, 11, "("),
                     token!(Ident, 6, 12, "a"),
                     token!(GreaterThan, 6, 14, ">"),
                     token!(IntLiteral, 6, 16, "20"),
                     token!(RParen, 6, 18, ")"),
                     token!(LBrace, 6, 20, "{"),
                     token!(Return, 7, 3, "return"),
                     token!(Minus, 7, 10, "-"),
                     token!(IntLiteral, 7, 11, "30"),
                     token!(Divide, 7, 14, "/"),
                     token!(IntLiteral, 7, 16, "40"),
                     token!(Multiply, 7, 19, "*"),
                     token!(IntLiteral, 7, 21, "50"),
                     token!(SemiColon, 7, 23, ";"),
                     token!(RBrace, 8, 1, "}"),
                     token!(Else, 8, 3, "else"),
                     token!(If, 8, 8, "if"),
                     token!(LParen, 8, 11, "("),
                     token!(Ident, 8, 12, "a"),
                     token!(LessThan, 8, 14, "<"),
                     token!(IntLiteral, 8, 16, "30"),
                     token!(RParen, 8, 18, ")"),
                     token!(LBrace, 8, 20, "{"),
                     token!(Return, 9, 3, "return"),
                     token!(BoolLiteral, 9, 10, "true"),
                     token!(SemiColon, 9, 14, ";"),
                     token!(RBrace, 10, 1, "}"),
                     token!(Return, 11, 1, "return"),
                     token!(BoolLiteral, 11, 8, "false"),
                     token!(SemiColon, 11, 13, ";")]);
    }

    #[test]
    fn id_with_numbers() {
        vec_eq(&lex("hello2 hel301oo120"),
               &vec![token!(Ident, 1, 1, "hello2"),
                     token!(Ident, 1, 8, "hel301oo120")]);
    }

    #[test]
    fn string_literals() {
        vec_eq(&lex("\"foobar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foobar\"")]);
        vec_eq(&lex("\"foo bar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo bar\"")]);
        vec_eq(&lex("\"foo\\nbar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\nbar\"")]);
        vec_eq(&lex("\"foo\\tbar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\tbar\"")]);
        vec_eq(&lex("\"foo\\\"bar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\\"bar\"")]);
    }

    #[test]
    fn array_tokens() {
        vec_eq(&lex("[1, 2];"),
               &vec![token!(LBracket, 1, 1, "["),
                     token!(IntLiteral, 1, 2, "1"),
                     token!(Comma, 1, 3, ","),
                     token!(IntLiteral, 1, 5, "2"),
                     token!(RBracket, 1, 6, "]"),
                     token!(SemiColon, 1, 7, ";")]);
    }

    #[test]
    fn hash_tokens() {
        vec_eq(&lex("{\"hello\": \"world\"}"),
               &vec![token!(LBrace, 1, 1, "{"),
                    token!(StringLiteral, 1, 2, "\"hello\""),
                    token!(Colon, 1, 9, ":"),
                    token!(StringLiteral, 1, 11, "\"world\""),
                    token!(RBrace, 1, 18, "}")]);
    }
}
<commit_msg>Implement lex_string()<commit_after>#[macro_use]
pub mod types;

use self::types::*;
use common::combinator::Parser;
use std::iter::FromIterator;

pub fn tokenize<T: Iterator<Item = char>>(input: T) -> Result<Vec<Token>> {
    let mut l = Lexer::new(input);

    let mut result = Vec::default();

    loop {
        let _ = skip_whitespaces(&mut l);
        match l.preview() {
            None => break,
            _ => {
                result.push(try!(lex_token(&mut l)));
            }
        }
    }

    Ok(result)
}

fn skip_whitespaces(l: &mut Lexer) -> Result<Vec<char>> {
    l.many(|l| l.predicate(|x| [' ', '\t', '\n', '\r'].contains(x)))
}

fn lex_token(l: &mut Lexer) -> Result<Token> {
    l.choose(&[&lex_operator, &lex_punctuation, &lex_string])
}

macro_rules! parse_map {
    ($pat:expr, $tok:ident) => {
        &|l| {
            let pos = l.current_pos();
            let lit: String = try!(l.string($pat.chars()));
            Ok(token!($tok, pos.0, pos.1, lit))
        }
    }
}

fn lex_operator(l: &mut Lexer) -> Result<Token> {
    l.choose(&[parse_map!("==", Eq),
               parse_map!("=", Assign),
               parse_map!("+", Plus),
               parse_map!("-", Minus),
               parse_map!("*", Multiply),
               parse_map!("/", Divide),
               parse_map!("!=", NotEq),
               parse_map!("!", Not),
               parse_map!(">", GreaterThan),
               parse_map!("<", LessThan)])
}

fn lex_punctuation(l: &mut Lexer) -> Result<Token> {
    l.choose(&[parse_map!(":", Colon),
               parse_map!(";", SemiColon),
               parse_map!(",", Comma),
               parse_map!("(", LParen),
               parse_map!(")", RParen),
               parse_map!("{", LBrace),
               parse_map!("}", RBrace),
               parse_map!("[", LBracket),
               parse_map!("]", RBracket)])
}

fn lex_string(l: &mut Lexer) -> Result<Token> {
    let pos = l.current_pos();
    let mut result = Vec::default();

    result.push(try!(l.atom('"')));
    loop {
        let c = try!(l.next());
        result.push(c);
        match c {
            '\\' => result.push(try!(l.next())),
            '"' => break,
            _ => continue,
        }
    }

    Ok(token!(StringLiteral, pos.0, pos.1, String::from_iter(result)))
}

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

    fn lex(input: &str) -> Vec<Token> {
        tokenize(String::from(input).chars()).unwrap()
    }

    fn vec_eq<T: PartialEq + Debug>(xs: &Vec<T>, ys: &Vec<T>) {
        if xs.len() != ys.len() {
            panic!("{:?} != {:?}", xs, ys);
        }
        for i in 0..xs.len() {
            if xs[i] != ys[i] {
                panic!("{:?} != {:?}", xs, ys);
            }
        }
    }

    #[test]
    fn special_chars() {
        vec_eq(&lex("=+(){},;"),
               &vec![token!(Assign, 1, 1, '='),
                     token!(Plus, 1, 2, '+'),
                     token!(LParen, 1, 3, '('),
                     token!(RParen, 1, 4, ')'),
                     token!(LBrace, 1, 5, '{'),
                     token!(RBrace, 1, 6, '}'),
                     token!(Comma, 1, 7, ','),
                     token!(SemiColon, 1, 8, ';')]);
    }

    #[test]
    fn complex_code() {
        vec_eq(&lex("
let five = 5;
let ten = 10;
let add = fn(x, y) {
    x + y;
};
let result = add(five, ten);
               "),
               &vec![token!(Let, 2, 1, "let"),
                     token!(Ident, 2, 5, "five"),
                     token!(Assign, 2, 10, "="),
                     token!(IntLiteral, 2, 12, "5"),
                     token!(SemiColon, 2, 13, ";"),
                     token!(Let, 3, 1, "let"),
                     token!(Ident, 3, 5, "ten"),
                     token!(Assign, 3, 9, "="),
                     token!(IntLiteral, 3, 11, "10"),
                     token!(SemiColon, 3, 13, ";"),
                     token!(Let, 4, 1, "let"),
                     token!(Ident, 4, 5, "add"),
                     token!(Assign, 4, 9, "="),
                     token!(Function, 4, 11, "fn"),
                     token!(LParen, 4, 13, "("),
                     token!(Ident, 4, 14, "x"),
                     token!(Comma, 4, 15, ","),
                     token!(Ident, 4, 17, "y"),
                     token!(RParen, 4, 18, ")"),
                     token!(LBrace, 4, 20, "{"),
                     token!(Ident, 5, 5, "x"),
                     token!(Plus, 5, 7, "+"),
                     token!(Ident, 5, 9, "y"),
                     token!(SemiColon, 5, 10, ";"),
                     token!(RBrace, 6, 1, "}"),
                     token!(SemiColon, 6, 2, ";"),
                     token!(Let, 7, 1, "let"),
                     token!(Ident, 7, 5, "result"),
                     token!(Assign, 7, 12, "="),
                     token!(Ident, 7, 14, "add"),
                     token!(LParen, 7, 17, "("),
                     token!(Ident, 7, 18, "five"),
                     token!(Comma, 7, 22, ","),
                     token!(Ident, 7, 24, "ten"),
                     token!(RParen, 7, 27, ")"),
                     token!(SemiColon, 7, 28, ";")]);
    }

    #[test]
    fn complex_code2() {
        vec_eq(&lex("
if (a == 10) {
  return a;
} else if (a != 20) {
  return !a;
} else if (a > 20) {
  return -30 / 40 * 50;
} else if (a < 30) {
  return true;
}
return false;
                    "),
               &vec![token!(If, 2, 1, "if"),
                     token!(LParen, 2, 3, "("),
                     token!(Ident, 2, 4, "a"),
                     token!(Eq, 2, 6, "=="),
                     token!(IntLiteral, 2, 9, "10"),
                     token!(RParen, 2, 11, ")"),
                     token!(LBrace, 2, 13, "{"),
                     token!(Return, 3, 3, "return"),
                     token!(Ident, 3, 10, "a"),
                     token!(SemiColon, 3, 11, ";"),
                     token!(RBrace, 4, 1, "}"),
                     token!(Else, 4, 3, "else"),
                     token!(If, 4, 8, "if"),
                     token!(LParen, 4, 11, "("),
                     token!(Ident, 4, 12, "a"),
                     token!(NotEq, 4, 14, "!="),
                     token!(IntLiteral, 4, 17, "20"),
                     token!(RParen, 4, 19, ")"),
                     token!(LBrace, 4, 21, "{"),
                     token!(Return, 5, 3, "return"),
                     token!(Not, 5, 10, "!"),
                     token!(Ident, 5, 11, "a"),
                     token!(SemiColon, 5, 13, ";"),
                     token!(RBrace, 6, 1, "}"),
                     token!(Else, 6, 3, "else"),
                     token!(If, 6, 8, "if"),
                     token!(LParen, 6, 11, "("),
                     token!(Ident, 6, 12, "a"),
                     token!(GreaterThan, 6, 14, ">"),
                     token!(IntLiteral, 6, 16, "20"),
                     token!(RParen, 6, 18, ")"),
                     token!(LBrace, 6, 20, "{"),
                     token!(Return, 7, 3, "return"),
                     token!(Minus, 7, 10, "-"),
                     token!(IntLiteral, 7, 11, "30"),
                     token!(Divide, 7, 14, "/"),
                     token!(IntLiteral, 7, 16, "40"),
                     token!(Multiply, 7, 19, "*"),
                     token!(IntLiteral, 7, 21, "50"),
                     token!(SemiColon, 7, 23, ";"),
                     token!(RBrace, 8, 1, "}"),
                     token!(Else, 8, 3, "else"),
                     token!(If, 8, 8, "if"),
                     token!(LParen, 8, 11, "("),
                     token!(Ident, 8, 12, "a"),
                     token!(LessThan, 8, 14, "<"),
                     token!(IntLiteral, 8, 16, "30"),
                     token!(RParen, 8, 18, ")"),
                     token!(LBrace, 8, 20, "{"),
                     token!(Return, 9, 3, "return"),
                     token!(BoolLiteral, 9, 10, "true"),
                     token!(SemiColon, 9, 14, ";"),
                     token!(RBrace, 10, 1, "}"),
                     token!(Return, 11, 1, "return"),
                     token!(BoolLiteral, 11, 8, "false"),
                     token!(SemiColon, 11, 13, ";")]);
    }

    #[test]
    fn id_with_numbers() {
        vec_eq(&lex("hello2 hel301oo120"),
               &vec![token!(Ident, 1, 1, "hello2"),
                     token!(Ident, 1, 8, "hel301oo120")]);
    }

    #[test]
    fn string_literals() {
        vec_eq(&lex("\"foobar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foobar\"")]);
        vec_eq(&lex("\"foo bar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo bar\"")]);
        vec_eq(&lex("\"foo\\nbar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\nbar\"")]);
        vec_eq(&lex("\"foo\\tbar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\tbar\"")]);
        vec_eq(&lex("\"foo\\\"bar\""),
               &vec![token!(StringLiteral, 1, 1, "\"foo\\\"bar\"")]);
    }

    #[test]
    fn array_tokens() {
        vec_eq(&lex("[1, 2];"),
               &vec![token!(LBracket, 1, 1, "["),
                     token!(IntLiteral, 1, 2, "1"),
                     token!(Comma, 1, 3, ","),
                     token!(IntLiteral, 1, 5, "2"),
                     token!(RBracket, 1, 6, "]"),
                     token!(SemiColon, 1, 7, ";")]);
    }

    #[test]
    fn hash_tokens() {
        vec_eq(&lex("{\"hello\": \"world\"}"),
               &vec![token!(LBrace, 1, 1, "{"),
                    token!(StringLiteral, 1, 2, "\"hello\""),
                    token!(Colon, 1, 9, ":"),
                    token!(StringLiteral, 1, 11, "\"world\""),
                    token!(RBrace, 1, 18, "}")]);
    }
}
<|endoftext|>


<commit_before>//! A simple wait-free, grow-only pagetable, assumes a dense keyspace.

use std::sync::atomic::Ordering::{Relaxed, SeqCst};

use super::*;

const FANFACTOR: usize = 18;
const FANOUT: usize = 1 << FANFACTOR;
const FAN_MASK: usize = FANOUT - 1;

pub type PageId = usize;

macro_rules! rep_no_copy {
    ($e:expr; $n:expr) => {{
        let mut v = Vec::with_capacity($n);

        for _ in 0..$n {
            v.push($e);
        }

        v
    }};
}

#[inline(always)]
fn split_fanout(i: usize) -> (usize, usize) {
    // right shift 32 on 32-bit pointer systems panics
    #[cfg(target_pointer_width = "64")]
    assert!(
        i <= 1 << (FANFACTOR * 2),
        "trying to access key of {}, which is \
         higher than 2 ^ {}",
        i,
        (FANFACTOR * 2)
    );

    let left = i >> FANFACTOR;
    let right = i & FAN_MASK;
    (left, right)
}

struct Node1<T: Send + 'static> {
    children: Vec<Atomic<Node2<T>>>,
}

struct Node2<T: Send + 'static> {
    children: Vec<Atomic<T>>,
}

impl<T: Send + 'static> Default for Node1<T> {
    fn default() -> Node1<T> {
        let children = rep_no_copy!(Atomic::null(); FANOUT);
        Node1 { children }
    }
}

impl<T: Send + 'static> Default for Node2<T> {
    fn default() -> Node2<T> {
        let children = rep_no_copy!(Atomic::null(); FANOUT);
        Node2 { children }
    }
}

/// A simple lock-free radix tree.
pub struct PageTable<T>
where
    T: 'static + Send + Sync,
{
    head: Atomic<Node1<T>>,
}

impl<T> Default for PageTable<T>
where
    T: 'static + Send + Sync,
{
    fn default() -> PageTable<T> {
        let head = Owned::new(Node1::default());
        PageTable {
            head: Atomic::from(head),
        }
    }
}

impl<T> PageTable<T>
where
    T: 'static + Send + Sync,
{
    /// Atomically swap the previous value in a tree with a new one.
    pub fn swap<'g>(
        &self,
        pid: PageId,
        new: Shared<'g, T>,
        guard: &'g Guard,
    ) -> Shared<'g, T> {
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);
        debug_delay();
        tip.swap(new, SeqCst, guard)
    }

    /// Compare and swap an old value to a new one.
    pub fn cas<'g>(
        &self,
        pid: PageId,
        old: Shared<'g, T>,
        new: Shared<'g, T>,
        guard: &'g Guard,
    ) -> std::result::Result<Shared<'g, T>, Shared<'g, T>> {
        debug_delay();
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);

        debug_delay();

        let _ = tip
            .compare_and_set(old, new, SeqCst, guard)
            .map_err(|e| e.current)?;

        if !old.is_null() {
            unsafe {
                guard.defer_destroy(old);
            }
        }
        Ok(new)
    }

    /// Try to get a value from the tree.
    pub fn get<'g>(
        &self,
        pid: PageId,
        guard: &'g Guard,
    ) -> Option<Shared<'g, T>> {
        debug_delay();
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);

        let res = tip.load(SeqCst, guard);
        if res.is_null() {
            None
        } else {
            Some(res)
        }
    }

    /// Delete a value from the tree, returning the old value if it was set.
    pub fn del<'g>(
        &self,
        pid: PageId,
        guard: &'g Guard,
    ) -> Option<Shared<'g, T>> {
        debug_delay();
        let old = self.swap(pid, Shared::null(), guard);
        if !old.is_null() {
            unsafe {
                guard.defer_destroy(old);
            }
            Some(old)
        } else {
            None
        }
    }
}

fn traverse<'g, T: 'static + Send>(
    head: Shared<'g, Node1<T>>,
    k: PageId,
    guard: &'g Guard,
) -> &'g Atomic<T> {
    let (l1k, l2k) = split_fanout(k);

    debug_delay();
    let l1 = unsafe { &head.deref().children };

    debug_delay();
    let mut l2_ptr = l1[l1k].load(SeqCst, guard);

    if l2_ptr.is_null() {
        let next_child =
            Owned::new(Node2::default()).into_shared(guard);

        debug_delay();
        let ret = l1[l1k]
            .compare_and_set(l2_ptr, next_child, SeqCst, guard);

        match ret {
            Ok(_) => {
                // CAS worked
                l2_ptr = next_child;
            }
            Err(e) => {
                // another thread beat us, drop unused created
                // child and use what is already set
                l2_ptr = e.current;
            }
        }
    }

    debug_delay();
    let l2 = unsafe { &l2_ptr.deref().children };

    &l2[l2k]
}

impl<T> Drop for PageTable<T>
where
    T: 'static + Send + Sync,
{
    fn drop(&mut self) {
        unsafe {
            let head =
                self.head.load(Relaxed, &unprotected()).as_raw()
                    as usize;
            drop(Box::from_raw(head as *mut Node1<T>));
        }
    }
}

impl<T: Send + 'static> Drop for Node1<T> {
    fn drop(&mut self) {
        unsafe {
            let children: Vec<*const Node2<T>> = self
                .children
                .iter()
                .map(|c| c.load(Relaxed, &unprotected()).as_raw())
                .filter(|c| !c.is_null())
                .collect();

            for child in children {
                drop(Box::from_raw(child as *mut Node2<T>));
            }
        }
    }
}

impl<T: Send + 'static> Drop for Node2<T> {
    fn drop(&mut self) {
        unsafe {
            let children: Vec<*const T> = self
                .children
                .iter()
                .map(|c| c.load(Relaxed, &unprotected()).as_raw())
                .filter(|c| !c.is_null())
                .collect();

            for child in children {
                drop(Box::from_raw(child as *mut T));
            }
        }
    }
}

#[test]
fn test_split_fanout() {
    assert_eq!(
        split_fanout(0b11_1111_1111_1111_1111),
        (0, 0b11_1111_1111_1111_1111)
    );
    assert_eq!(
        split_fanout(0b111_1111_1111_1111_1111),
        (0b1, 0b11_1111_1111_1111_1111)
    );
}

#[test]
fn basic_functionality() {
    unsafe {
        let guard = pin();
        let rt = PageTable::default();
        let v1 = Owned::new(5).into_shared(&guard);
        rt.cas(0, Shared::null(), v1, &guard).unwrap();
        let ptr = rt.get(0, &guard).unwrap();
        assert_eq!(ptr.deref(), &5);
        rt.cas(0, ptr, Owned::new(6).into_shared(&guard), &guard)
            .unwrap();
        assert_eq!(rt.get(0, &guard).unwrap().deref(), &6);
        rt.del(0, &guard);
        assert!(rt.get(0, &guard).is_none());

        let k2 = 321 << FANFACTOR;
        let k3 = k2 + 1;

        let v2 = Owned::new(2).into_shared(&guard);
        rt.cas(k2, Shared::null(), v2, &guard).unwrap();
        assert_eq!(rt.get(k2, &guard).unwrap().deref(), &2);
        assert!(rt.get(k3, &guard).is_none());
        let v3 = Owned::new(3).into_shared(&guard);
        rt.cas(k3, Shared::null(), v3, &guard).unwrap();
        assert_eq!(rt.get(k3, &guard).unwrap().deref(), &3);
        assert_eq!(rt.get(k2, &guard).unwrap().deref(), &2);
    }
}

#[test]
fn test_model() {
    use self::Shared as EpochShared;
    use model::{model, prop_oneof};

    model! {
        Model => let mut m = std::collections::HashMap::new(),
        Implementation => let i = PageTable::default(),
        Insert((usize, usize))((k, new) in (0usize..4, 0usize..4)) => {
            if !m.contains_key(&k) {
                m.insert(k, new);

                let guard = pin();
                let v = Owned::new(new).into_shared(&guard);
                i.cas(k, EpochShared::null(), v, &guard ).expect("should be able to insert a value");
            }
        },
        Get(usize)(k in 0usize..4) => {
            let guard = pin();
            let expected = m.get(&k);
            let actual = i.get(k, &guard).map(|s| unsafe { s.deref() });
            assert_eq!(expected, actual);
        },
        Cas((usize, usize, usize))((k, old, new) in (0usize..4, 0usize..4, 0usize..4)) => {
            let guard = pin();
            let expected_current = m.get(&k).cloned();
            let actual_current = i.get(k, &guard);
            assert_eq!(expected_current, actual_current.map(|s| unsafe { *s.deref() }));
            if expected_current.is_none() {
                continue;
            }

            let new_v = Owned::new(new).into_shared(&guard);

            if expected_current == Some(old) {
                m.insert(k, new);
                let cas_res = i.cas(k, actual_current.unwrap(), new_v, &guard);
                assert!(cas_res.is_ok());
            };
        }
    }
}

#[test]
#[ignore]
fn test_linearizability() {
    use self::Shared as EpochShared;
    use model::{linearizable, prop_oneof, Shared};

    linearizable! {
        Implementation => let i = Shared::new(PageTable::default()),
        Get(usize)(k in 0usize..4) -> Option<usize> {
            let guard = pin();
            unsafe {
                i.get(k, &guard).map(|s| *s.deref())
            }
        },
        Insert((usize, usize))((k, new) in (0usize..4, 0usize..4)) -> bool {
            let guard = pin();
            let v = Owned::new(new).into_shared(&guard);
            i.cas(k, EpochShared::null(), v, &guard).is_err()
        },
        Cas((usize, usize, usize))((k, old, new)
            in (0usize..4, 0usize..4, 0usize..4))
            -> std::result::Result<(), usize> {
            let guard = pin();
            i.cas(k, Owned::new(old).into_shared(&guard), Owned::new(new).into_shared(&guard), &guard)
                .map(|_| ())
                .map_err(|s| if s.is_null() {
                    0
                } else {
                    unsafe { *s.deref() }
                })
        }
    }
}
<commit_msg>Disable pagetable model check for now<commit_after>//! A simple wait-free, grow-only pagetable, assumes a dense keyspace.

use std::sync::atomic::Ordering::{Relaxed, SeqCst};

use super::*;

const FANFACTOR: usize = 18;
const FANOUT: usize = 1 << FANFACTOR;
const FAN_MASK: usize = FANOUT - 1;

pub type PageId = usize;

macro_rules! rep_no_copy {
    ($e:expr; $n:expr) => {{
        let mut v = Vec::with_capacity($n);

        for _ in 0..$n {
            v.push($e);
        }

        v
    }};
}

#[inline(always)]
fn split_fanout(i: usize) -> (usize, usize) {
    // right shift 32 on 32-bit pointer systems panics
    #[cfg(target_pointer_width = "64")]
    assert!(
        i <= 1 << (FANFACTOR * 2),
        "trying to access key of {}, which is \
         higher than 2 ^ {}",
        i,
        (FANFACTOR * 2)
    );

    let left = i >> FANFACTOR;
    let right = i & FAN_MASK;
    (left, right)
}

struct Node1<T: Send + 'static> {
    children: Vec<Atomic<Node2<T>>>,
}

struct Node2<T: Send + 'static> {
    children: Vec<Atomic<T>>,
}

impl<T: Send + 'static> Default for Node1<T> {
    fn default() -> Node1<T> {
        let children = rep_no_copy!(Atomic::null(); FANOUT);
        Node1 { children }
    }
}

impl<T: Send + 'static> Default for Node2<T> {
    fn default() -> Node2<T> {
        let children = rep_no_copy!(Atomic::null(); FANOUT);
        Node2 { children }
    }
}

/// A simple lock-free radix tree.
pub struct PageTable<T>
where
    T: 'static + Send + Sync,
{
    head: Atomic<Node1<T>>,
}

impl<T> Default for PageTable<T>
where
    T: 'static + Send + Sync,
{
    fn default() -> PageTable<T> {
        let head = Owned::new(Node1::default());
        PageTable {
            head: Atomic::from(head),
        }
    }
}

impl<T> PageTable<T>
where
    T: 'static + Send + Sync,
{
    /// Atomically swap the previous value in a tree with a new one.
    pub fn swap<'g>(
        &self,
        pid: PageId,
        new: Shared<'g, T>,
        guard: &'g Guard,
    ) -> Shared<'g, T> {
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);
        debug_delay();
        tip.swap(new, SeqCst, guard)
    }

    /// Compare and swap an old value to a new one.
    pub fn cas<'g>(
        &self,
        pid: PageId,
        old: Shared<'g, T>,
        new: Shared<'g, T>,
        guard: &'g Guard,
    ) -> std::result::Result<Shared<'g, T>, Shared<'g, T>> {
        debug_delay();
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);

        debug_delay();

        let _ = tip
            .compare_and_set(old, new, SeqCst, guard)
            .map_err(|e| e.current)?;

        if !old.is_null() {
            unsafe {
                guard.defer_destroy(old);
            }
        }
        Ok(new)
    }

    /// Try to get a value from the tree.
    pub fn get<'g>(
        &self,
        pid: PageId,
        guard: &'g Guard,
    ) -> Option<Shared<'g, T>> {
        debug_delay();
        let tip = traverse(self.head.load(SeqCst, guard), pid, guard);

        let res = tip.load(SeqCst, guard);
        if res.is_null() {
            None
        } else {
            Some(res)
        }
    }

    /// Delete a value from the tree, returning the old value if it was set.
    pub fn del<'g>(
        &self,
        pid: PageId,
        guard: &'g Guard,
    ) -> Option<Shared<'g, T>> {
        debug_delay();
        let old = self.swap(pid, Shared::null(), guard);
        if !old.is_null() {
            unsafe {
                guard.defer_destroy(old);
            }
            Some(old)
        } else {
            None
        }
    }
}

fn traverse<'g, T: 'static + Send>(
    head: Shared<'g, Node1<T>>,
    k: PageId,
    guard: &'g Guard,
) -> &'g Atomic<T> {
    let (l1k, l2k) = split_fanout(k);

    debug_delay();
    let l1 = unsafe { &head.deref().children };

    debug_delay();
    let mut l2_ptr = l1[l1k].load(SeqCst, guard);

    if l2_ptr.is_null() {
        let next_child =
            Owned::new(Node2::default()).into_shared(guard);

        debug_delay();
        let ret = l1[l1k]
            .compare_and_set(l2_ptr, next_child, SeqCst, guard);

        match ret {
            Ok(_) => {
                // CAS worked
                l2_ptr = next_child;
            }
            Err(e) => {
                // another thread beat us, drop unused created
                // child and use what is already set
                l2_ptr = e.current;
            }
        }
    }

    debug_delay();
    let l2 = unsafe { &l2_ptr.deref().children };

    &l2[l2k]
}

impl<T> Drop for PageTable<T>
where
    T: 'static + Send + Sync,
{
    fn drop(&mut self) {
        unsafe {
            let head =
                self.head.load(Relaxed, &unprotected()).as_raw()
                    as usize;
            drop(Box::from_raw(head as *mut Node1<T>));
        }
    }
}

impl<T: Send + 'static> Drop for Node1<T> {
    fn drop(&mut self) {
        unsafe {
            let children: Vec<*const Node2<T>> = self
                .children
                .iter()
                .map(|c| c.load(Relaxed, &unprotected()).as_raw())
                .filter(|c| !c.is_null())
                .collect();

            for child in children {
                drop(Box::from_raw(child as *mut Node2<T>));
            }
        }
    }
}

impl<T: Send + 'static> Drop for Node2<T> {
    fn drop(&mut self) {
        unsafe {
            let children: Vec<*const T> = self
                .children
                .iter()
                .map(|c| c.load(Relaxed, &unprotected()).as_raw())
                .filter(|c| !c.is_null())
                .collect();

            for child in children {
                drop(Box::from_raw(child as *mut T));
            }
        }
    }
}

#[test]
fn test_split_fanout() {
    assert_eq!(
        split_fanout(0b11_1111_1111_1111_1111),
        (0, 0b11_1111_1111_1111_1111)
    );
    assert_eq!(
        split_fanout(0b111_1111_1111_1111_1111),
        (0b1, 0b11_1111_1111_1111_1111)
    );
}

#[test]
fn basic_functionality() {
    unsafe {
        let guard = pin();
        let rt = PageTable::default();
        let v1 = Owned::new(5).into_shared(&guard);
        rt.cas(0, Shared::null(), v1, &guard).unwrap();
        let ptr = rt.get(0, &guard).unwrap();
        assert_eq!(ptr.deref(), &5);
        rt.cas(0, ptr, Owned::new(6).into_shared(&guard), &guard)
            .unwrap();
        assert_eq!(rt.get(0, &guard).unwrap().deref(), &6);
        rt.del(0, &guard);
        assert!(rt.get(0, &guard).is_none());

        let k2 = 321 << FANFACTOR;
        let k3 = k2 + 1;

        let v2 = Owned::new(2).into_shared(&guard);
        rt.cas(k2, Shared::null(), v2, &guard).unwrap();
        assert_eq!(rt.get(k2, &guard).unwrap().deref(), &2);
        assert!(rt.get(k3, &guard).is_none());
        let v3 = Owned::new(3).into_shared(&guard);
        rt.cas(k3, Shared::null(), v3, &guard).unwrap();
        assert_eq!(rt.get(k3, &guard).unwrap().deref(), &3);
        assert_eq!(rt.get(k2, &guard).unwrap().deref(), &2);
    }
}

#[test]
#[ignore]
fn test_model() {
    use self::Shared as EpochShared;
    use model::{model, prop_oneof};

    model! {
        Model => let mut m = std::collections::HashMap::new(),
        Implementation => let i = PageTable::default(),
        Insert((usize, usize))((k, new) in (0usize..4, 0usize..4)) => {
            if !m.contains_key(&k) {
                m.insert(k, new);

                let guard = pin();
                let v = Owned::new(new).into_shared(&guard);
                i.cas(k, EpochShared::null(), v, &guard ).expect("should be able to insert a value");
            }
        },
        Get(usize)(k in 0usize..4) => {
            let guard = pin();
            let expected = m.get(&k);
            let actual = i.get(k, &guard).map(|s| unsafe { s.deref() });
            assert_eq!(expected, actual);
        },
        Cas((usize, usize, usize))((k, old, new) in (0usize..4, 0usize..4, 0usize..4)) => {
            let guard = pin();
            let expected_current = m.get(&k).cloned();
            let actual_current = i.get(k, &guard);
            assert_eq!(expected_current, actual_current.map(|s| unsafe { *s.deref() }));
            if expected_current.is_none() {
                continue;
            }

            let new_v = Owned::new(new).into_shared(&guard);

            if expected_current == Some(old) {
                m.insert(k, new);
                let cas_res = i.cas(k, actual_current.unwrap(), new_v, &guard);
                assert!(cas_res.is_ok());
            };
        }
    }
}

#[test]
#[ignore]
fn test_linearizability() {
    use self::Shared as EpochShared;
    use model::{linearizable, prop_oneof, Shared};

    linearizable! {
        Implementation => let i = Shared::new(PageTable::default()),
        Get(usize)(k in 0usize..4) -> Option<usize> {
            let guard = pin();
            unsafe {
                i.get(k, &guard).map(|s| *s.deref())
            }
        },
        Insert((usize, usize))((k, new) in (0usize..4, 0usize..4)) -> bool {
            let guard = pin();
            let v = Owned::new(new).into_shared(&guard);
            i.cas(k, EpochShared::null(), v, &guard).is_err()
        },
        Cas((usize, usize, usize))((k, old, new)
            in (0usize..4, 0usize..4, 0usize..4))
            -> std::result::Result<(), usize> {
            let guard = pin();
            i.cas(k, Owned::new(old).into_shared(&guard), Owned::new(new).into_shared(&guard), &guard)
                .map(|_| ())
                .map_err(|s| if s.is_null() {
                    0
                } else {
                    unsafe { *s.deref() }
                })
        }
    }
}
<|endoftext|>