<commit_before>use ring::aead::{open_in_place, seal_in_place, Aad, Algorithm, Nonce, AES_256_GCM};
use ring::aead::{OpeningKey, SealingKey};
use ring::rand::{SecureRandom, SystemRandom};

use super::Key;
use crate::cookie::{Cookie, CookieJar};

// Keep these in sync, and keep the key len synced with the `private` docs as
// well as the `KEYS_INFO` const in secure::Key.
static ALGO: &'static Algorithm = &AES_256_GCM;
const NONCE_LEN: usize = 12;
pub const KEY_LEN: usize = 32;

/// A child cookie jar that provides authenticated encryption for its cookies.
///
/// A _private_ child jar signs and encrypts all the cookies added to it and
/// verifies and decrypts cookies retrieved from it. Any cookies stored in a
/// `PrivateJar` are simultaneously assured confidentiality, integrity, and
/// authenticity. In other words, clients cannot discover nor tamper with the
/// contents of a cookie, nor can they fabricate cookie data.
///
/// This type is only available when the `secure` feature is enabled.
pub struct PrivateJar<'a> {
    parent: &'a mut CookieJar,
    key: [u8; KEY_LEN],
}

impl<'a> PrivateJar<'a> {
    /// Creates a new child `PrivateJar` with parent `parent` and key `key`.
    /// This method is typically called indirectly via the `signed` method of
    /// `CookieJar`.
    #[doc(hidden)]
    pub fn new(parent: &'a mut CookieJar, key: &Key) -> PrivateJar<'a> {
        let mut key_array = [0u8; KEY_LEN];
        key_array.copy_from_slice(key.encryption());
        PrivateJar {
            parent: parent,
            key: key_array,
        }
    }

    /// Given a sealed value `str` and a key name `name`, where the nonce is
    /// prepended to the original value and then both are Base64 encoded,
    /// verifies and decrypts the sealed value and returns it. If there's a
    /// problem, returns an `Err` with a string describing the issue.
    fn unseal(&self, name: &str, value: &str) -> Result<String, &'static str> {
        let mut data = base64::decode(value).map_err(|_| "bad base64 value")?;
        if data.len() <= NONCE_LEN {
            return Err("length of decoded data is <= NONCE_LEN");
        }

        let ad = Aad::from(name.as_bytes());
        let key = OpeningKey::new(ALGO, &self.key).expect("opening key");
        let (nonce, sealed) = data.split_at_mut(NONCE_LEN);
        let nonce =
            Nonce::try_assume_unique_for_key(nonce).expect("invalid length of `nonce`");
        let unsealed = open_in_place(&key, nonce, ad, 0, sealed)
            .map_err(|_| "invalid key/nonce/value: bad seal")?;

        ::std::str::from_utf8(unsealed)
            .map(|s| s.to_string())
            .map_err(|_| "bad unsealed utf8")
    }

    /// Returns a reference to the `Cookie` inside this jar with the name `name`
    /// and authenticates and decrypts the cookie's value, returning a `Cookie`
    /// with the decrypted value. If the cookie cannot be found, or the cookie
    /// fails to authenticate or decrypt, `None` is returned.
    ///
    /// # Example
    ///
    /// ```rust
    /// use actix_http::cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// let mut private_jar = jar.private(&key);
    /// assert!(private_jar.get("name").is_none());
    ///
    /// private_jar.add(Cookie::new("name", "value"));
    /// assert_eq!(private_jar.get("name").unwrap().value(), "value");
    /// ```
    pub fn get(&self, name: &str) -> Option<Cookie<'static>> {
        if let Some(cookie_ref) = self.parent.get(name) {
            let mut cookie = cookie_ref.clone();
            if let Ok(value) = self.unseal(name, cookie.value()) {
                cookie.set_value(value);
                return Some(cookie);
            }
        }

        None
    }

    /// Adds `cookie` to the parent jar. The cookie's value is encrypted with
    /// authenticated encryption assuring confidentiality, integrity, and
    /// authenticity.
    ///
    /// # Example
    ///
    /// ```rust
    /// use actix_http::cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// jar.private(&key).add(Cookie::new("name", "value"));
    ///
    /// assert_ne!(jar.get("name").unwrap().value(), "value");
    /// assert_eq!(jar.private(&key).get("name").unwrap().value(), "value");
    /// ```
    pub fn add(&mut self, mut cookie: Cookie<'static>) {
        self.encrypt_cookie(&mut cookie);

        // Add the sealed cookie to the parent.
        self.parent.add(cookie);
    }

    /// Adds an "original" `cookie` to parent jar. The cookie's value is
    /// encrypted with authenticated encryption assuring confidentiality,
    /// integrity, and authenticity. Adding an original cookie does not affect
    /// the [`CookieJar::delta()`](struct.CookieJar.html#method.delta)
    /// computation. This method is intended to be used to seed the cookie jar
    /// with cookies received from a client's HTTP message.
    ///
    /// For accurate `delta` computations, this method should not be called
    /// after calling `remove`.
    ///
    /// # Example
    ///
    /// ```rust
    /// use actix_http::cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// jar.private(&key).add_original(Cookie::new("name", "value"));
    ///
    /// assert_eq!(jar.iter().count(), 1);
    /// assert_eq!(jar.delta().count(), 0);
    /// ```
    pub fn add_original(&mut self, mut cookie: Cookie<'static>) {
        self.encrypt_cookie(&mut cookie);

        // Add the sealed cookie to the parent.
        self.parent.add_original(cookie);
    }

    /// Encrypts the cookie's value with
    /// authenticated encryption assuring confidentiality, integrity, and authenticity.
    fn encrypt_cookie(&self, cookie: &mut Cookie) {
        let mut data;
        let output_len = {
            // Create the `SealingKey` structure.
            let key = SealingKey::new(ALGO, &self.key).expect("sealing key creation");

            // Create a vec to hold the [nonce | cookie value | overhead].
            let overhead = ALGO.tag_len();
            let cookie_val = cookie.value().as_bytes();
            data = vec![0; NONCE_LEN + cookie_val.len() + overhead];

            // Randomly generate the nonce, then copy the cookie value as input.
            let (nonce, in_out) = data.split_at_mut(NONCE_LEN);
            SystemRandom::new()
                .fill(nonce)
                .expect("couldn't random fill nonce");
            in_out[..cookie_val.len()].copy_from_slice(cookie_val);
            let nonce = Nonce::try_assume_unique_for_key(nonce)
                .expect("invalid length of `nonce`");

            // Use cookie's name as associated data to prevent value swapping.
            let ad = Aad::from(cookie.name().as_bytes());

            // Perform the actual sealing operation and get the output length.
            seal_in_place(&key, nonce, ad, in_out, overhead).expect("in-place seal")
        };

        // Base64 encode the nonce and encrypted value.
        let sealed_value = base64::encode(&data[..(NONCE_LEN + output_len)]);
        cookie.set_value(sealed_value);
    }

    /// Removes `cookie` from the parent jar.
    ///
    /// For correct removal, the passed in `cookie` must contain the same `path`
    /// and `domain` as the cookie that was initially set.
    ///
    /// See [CookieJar::remove](struct.CookieJar.html#method.remove) for more
    /// details.
    ///
    /// # Example
    ///
    /// ```rust
    /// use actix_http::cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// let mut private_jar = jar.private(&key);
    ///
    /// private_jar.add(Cookie::new("name", "value"));
    /// assert!(private_jar.get("name").is_some());
    ///
    /// private_jar.remove(Cookie::named("name"));
    /// assert!(private_jar.get("name").is_none());
    /// ```
    pub fn remove(&mut self, cookie: Cookie<'static>) {
        self.parent.remove(cookie);
    }
}

#[cfg(test)]
mod test {
    use super::{Cookie, CookieJar, Key};

    #[test]
    fn simple() {
        let key = Key::generate();
        let mut jar = CookieJar::new();
        assert_simple_behaviour!(jar, jar.private(&key));
    }

    #[test]
    fn private() {
        let key = Key::generate();
        let mut jar = CookieJar::new();
        assert_secure_behaviour!(jar, jar.private(&key));
    }
}
<commit_msg>Warn when an unsealed private cookie isn't valid UTF-8 (#746)<commit_after>use std::str;

use log::warn;
use ring::aead::{open_in_place, seal_in_place, Aad, Algorithm, Nonce, AES_256_GCM};
use ring::aead::{OpeningKey, SealingKey};
use ring::rand::{SecureRandom, SystemRandom};

use super::Key;
use crate::cookie::{Cookie, CookieJar};

// Keep these in sync, and keep the key len synced with the `private` docs as
// well as the `KEYS_INFO` const in secure::Key.
static ALGO: &'static Algorithm = &AES_256_GCM;
const NONCE_LEN: usize = 12;
pub const KEY_LEN: usize = 32;

/// A child cookie jar that provides authenticated encryption for its cookies.
///
/// A _private_ child jar signs and encrypts all the cookies added to it and
/// verifies and decrypts cookies retrieved from it. Any cookies stored in a
/// `PrivateJar` are simultaneously assured confidentiality, integrity, and
/// authenticity. In other words, clients cannot discover nor tamper with the
/// contents of a cookie, nor can they fabricate cookie data.
///
/// This type is only available when the `secure` feature is enabled.
pub struct PrivateJar<'a> {
    parent: &'a mut CookieJar,
    key: [u8; KEY_LEN],
}

impl<'a> PrivateJar<'a> {
    /// Creates a new child `PrivateJar` with parent `parent` and key `key`.
    /// This method is typically called indirectly via the `signed` method of
    /// `CookieJar`.
    #[doc(hidden)]
    pub fn new(parent: &'a mut CookieJar, key: &Key) -> PrivateJar<'a> {
        let mut key_array = [0u8; KEY_LEN];
        key_array.copy_from_slice(key.encryption());
        PrivateJar {
            parent: parent,
            key: key_array,
        }
    }

    /// Given a sealed value `str` and a key name `name`, where the nonce is
    /// prepended to the original value and then both are Base64 encoded,
    /// verifies and decrypts the sealed value and returns it. If there's a
    /// problem, returns an `Err` with a string describing the issue.
    fn unseal(&self, name: &str, value: &str) -> Result<String, &'static str> {
        let mut data = base64::decode(value).map_err(|_| "bad base64 value")?;
        if data.len() <= NONCE_LEN {
            return Err("length of decoded data is <= NONCE_LEN");
        }

        let ad = Aad::from(name.as_bytes());
        let key = OpeningKey::new(ALGO, &self.key).expect("opening key");
        let (nonce, sealed) = data.split_at_mut(NONCE_LEN);
        let nonce =
            Nonce::try_assume_unique_for_key(nonce).expect("invalid length of `nonce`");
        let unsealed = open_in_place(&key, nonce, ad, 0, sealed)
            .map_err(|_| "invalid key/nonce/value: bad seal")?;

        if let Ok(unsealed_utf8) = str::from_utf8(unsealed) {
            Ok(unsealed_utf8.to_string())
        } else {
            warn!("Private cookie does not have utf8 content!
It is likely the secret key used to encrypt them has been leaked.
Please change it as soon as possible.");
            Err("bad unsealed utf8")
        }
    }

    /// Returns a reference to the `Cookie` inside this jar with the name `name`
    /// and authenticates and decrypts the cookie's value, returning a `Cookie`
    /// with the decrypted value. If the cookie cannot be found, or the cookie
    /// fails to authenticate or decrypt, `None` is returned.
    ///
    /// # Example
    ///
    /// ```rust
    /// use actix_http::cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// let mut private_jar = jar.private(&key);
    /// assert!(private_jar.get("name").is_none());
    ///
    /// private_jar.add(Cookie::new("name", "value"));
    /// assert_eq!(private_jar.get("name").unwrap().value(), "value");
    /// ```
    pub fn get(&self, name: &str) -> Option<Cookie<'static>> {
        if let Some(cookie_ref) = self.parent.get(name) {
            let mut cookie = cookie_ref.clone();
            if let Ok(value) = self.unseal(name, cookie.value()) {
                cookie.set_value(value);
                return Some(cookie);
            }
        }

        None
    }

    /// Adds `cookie` to the parent jar. The cookie's value is encrypted with
    /// authenticated encryption assuring confidentiality, integrity, and
    /// authenticity.
    ///
    /// # Example
    ///
    /// ```rust
    /// use actix_http::cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// jar.private(&key).add(Cookie::new("name", "value"));
    ///
    /// assert_ne!(jar.get("name").unwrap().value(), "value");
    /// assert_eq!(jar.private(&key).get("name").unwrap().value(), "value");
    /// ```
    pub fn add(&mut self, mut cookie: Cookie<'static>) {
        self.encrypt_cookie(&mut cookie);

        // Add the sealed cookie to the parent.
        self.parent.add(cookie);
    }

    /// Adds an "original" `cookie` to parent jar. The cookie's value is
    /// encrypted with authenticated encryption assuring confidentiality,
    /// integrity, and authenticity. Adding an original cookie does not affect
    /// the [`CookieJar::delta()`](struct.CookieJar.html#method.delta)
    /// computation. This method is intended to be used to seed the cookie jar
    /// with cookies received from a client's HTTP message.
    ///
    /// For accurate `delta` computations, this method should not be called
    /// after calling `remove`.
    ///
    /// # Example
    ///
    /// ```rust
    /// use actix_http::cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// jar.private(&key).add_original(Cookie::new("name", "value"));
    ///
    /// assert_eq!(jar.iter().count(), 1);
    /// assert_eq!(jar.delta().count(), 0);
    /// ```
    pub fn add_original(&mut self, mut cookie: Cookie<'static>) {
        self.encrypt_cookie(&mut cookie);

        // Add the sealed cookie to the parent.
        self.parent.add_original(cookie);
    }

    /// Encrypts the cookie's value with
    /// authenticated encryption assuring confidentiality, integrity, and authenticity.
    fn encrypt_cookie(&self, cookie: &mut Cookie) {
        let name = cookie.name().as_bytes();
        let value = cookie.value().as_bytes();
        let data = encrypt_name_value(name, value, &self.key);

        // Base64 encode the nonce and encrypted value.
        let sealed_value = base64::encode(&data);
        cookie.set_value(sealed_value);
    }

    /// Removes `cookie` from the parent jar.
    ///
    /// For correct removal, the passed in `cookie` must contain the same `path`
    /// and `domain` as the cookie that was initially set.
    ///
    /// See [CookieJar::remove](struct.CookieJar.html#method.remove) for more
    /// details.
    ///
    /// # Example
    ///
    /// ```rust
    /// use actix_http::cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// let mut private_jar = jar.private(&key);
    ///
    /// private_jar.add(Cookie::new("name", "value"));
    /// assert!(private_jar.get("name").is_some());
    ///
    /// private_jar.remove(Cookie::named("name"));
    /// assert!(private_jar.get("name").is_none());
    /// ```
    pub fn remove(&mut self, cookie: Cookie<'static>) {
        self.parent.remove(cookie);
    }
}

fn encrypt_name_value(name: &[u8], value: &[u8], key: &[u8]) -> Vec<u8> {
    // Create the `SealingKey` structure.
    let key = SealingKey::new(ALGO, key).expect("sealing key creation");

    // Create a vec to hold the [nonce | cookie value | overhead].
    let overhead = ALGO.tag_len();
    let mut data = vec![0; NONCE_LEN + value.len() + overhead];

    // Randomly generate the nonce, then copy the cookie value as input.
    let (nonce, in_out) = data.split_at_mut(NONCE_LEN);
    SystemRandom::new()
        .fill(nonce)
        .expect("couldn't random fill nonce");
    in_out[..value.len()].copy_from_slice(value);
    let nonce = Nonce::try_assume_unique_for_key(nonce)
        .expect("invalid length of `nonce`");

    // Use cookie's name as associated data to prevent value swapping.
    let ad = Aad::from(name);

    // Perform the actual sealing operation and get the output length.
    let output_len = seal_in_place(&key, nonce, ad, in_out, overhead)
        .expect("in-place seal");

    // Remove the overhead and return the sealed content.
    data.truncate(NONCE_LEN + output_len);
    data
}

#[cfg(test)]
mod test {
    use super::{Cookie, CookieJar, Key, encrypt_name_value};

    #[test]
    fn simple() {
        let key = Key::generate();
        let mut jar = CookieJar::new();
        assert_simple_behaviour!(jar, jar.private(&key));
    }

    #[test]
    fn private() {
        let key = Key::generate();
        let mut jar = CookieJar::new();
        assert_secure_behaviour!(jar, jar.private(&key));
    }

    #[test]
    fn non_utf8() {
        let key = Key::generate();
        let mut jar = CookieJar::new();

        let name = "malicious";
        let mut assert_non_utf8 = |value: &[u8]| {
            let sealed = encrypt_name_value(name.as_bytes(), value, &key.encryption());
            let encoded = base64::encode(&sealed);
            assert_eq!(jar.private(&key).unseal(name, &encoded), Err("bad unsealed utf8"));
            jar.add(Cookie::new(name, encoded));
            assert_eq!(jar.private(&key).get(name), None);
        };

        assert_non_utf8(&[0x72, 0xfb, 0xdf, 0x74]);  // rûst in ISO/IEC 8859-1

        let mut malicious = String::from(r#"{"id":"abc123??%X","admin":true}"#)
            .into_bytes();
        malicious[8] |= 0b1100_0000;
        malicious[9] |= 0b1100_0000;
        assert_non_utf8(&malicious);
    }
}
<|endoftext|>


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

use crate::{FuzzTarget, FuzzTargetImpl};
use anyhow::{format_err, Result};
use once_cell::sync::Lazy;
use std::{collections::BTreeMap, env};

// List fuzz target modules here.
mod consensus;
mod executor;
mod json_rpc_service;
mod mempool;
mod move_vm;
mod network;
mod proof;
mod safety_rules;
mod secure_storage_vault;
mod state_sync;
mod storage;
mod transaction;
mod vm;

static ALL_TARGETS: Lazy<BTreeMap<&'static str, Box<dyn FuzzTargetImpl>>> = Lazy::new(|| {
    // List fuzz targets here in this format:
    let targets: Vec<Box<dyn FuzzTargetImpl>> = vec![
        // Consensus
        Box::new(consensus::ConsensusProposal::default()),
        // Executor
        Box::new(executor::ExecuteAndCommitBlocks::default()),
        Box::new(executor::ExecuteAndCommitChunk::default()),
        // JSON RPC Service
        Box::new(json_rpc_service::JsonRpcSubmitTransactionRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetAccountRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetAccountTransactionRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetAccountTransactionsRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetTransactionsRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetEventsRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetMetadataRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetCurrenciesRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetStateProofRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetAccountStateWithProofRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetNetworkStatusRequest::default()),
        // Mempool
        Box::new(mempool::MempoolIncomingTransactions::default()),
        // Move VM
        Box::new(move_vm::ValueTarget::default()),
        // Proof
        Box::new(proof::TestAccumulatorProofFuzzer::default()),
        Box::new(proof::SparseMerkleProofFuzzer::default()),
        Box::new(proof::TestAccumulatorRangeProofFuzzer::default()),
        Box::new(proof::TransactionInfoWithProofFuzzer::default()),
        Box::new(proof::AccountStateProofFuzzer::default()),
        Box::new(proof::EventProofFuzzer::default()),
        Box::new(proof::TransactionListProofFuzzer::default()),
        // Network
        Box::new(network::NetworkNoiseInitiator::default()),
        Box::new(network::NetworkNoiseResponder::default()),
        Box::new(network::NetworkNoiseStream::default()),
        Box::new(network::NetworkHandshakeExchange::default()),
        Box::new(network::NetworkHandshakeNegotiation::default()),
        Box::new(network::PeerNetworkMessagesReceive::default()),
        // Safety Rules Server (LSR)
        Box::new(safety_rules::SafetyRulesConstructAndSignVote::default()),
        Box::new(safety_rules::SafetyRulesInitialize::default()),
        Box::new(safety_rules::SafetyRulesHandleMessage::default()),
        Box::new(safety_rules::SafetyRulesSignProposal::default()),
        Box::new(safety_rules::SafetyRulesSignTimeout::default()),
        // Secure Storage Vault
        Box::new(secure_storage_vault::VaultGenericResponse::default()),
        Box::new(secure_storage_vault::VaultPolicyReadResponse::default()),
        Box::new(secure_storage_vault::VaultPolicyListResponse::default()),
        Box::new(secure_storage_vault::VaultSecretListResponse::default()),
        Box::new(secure_storage_vault::VaultSecretReadResponse::default()),
        Box::new(secure_storage_vault::VaultTokenCreateResponse::default()),
        Box::new(secure_storage_vault::VaultTokenRenewResponse::default()),
        Box::new(secure_storage_vault::VaultTransitCreateResponse::default()),
        Box::new(secure_storage_vault::VaultTransitExportResponse::default()),
        Box::new(secure_storage_vault::VaultTransitListResponse::default()),
        Box::new(secure_storage_vault::VaultTransitReadResponse::default()),
        Box::new(secure_storage_vault::VaultTransitRestoreResponse::default()),
        Box::new(secure_storage_vault::VaultTransitSignResponse::default()),
        Box::new(secure_storage_vault::VaultUnsealedResponse::default()),
        // State Sync
        Box::new(state_sync::StateSyncMsg::default()),
        // Storage
        // Box::new(storage::StorageSaveBlocks::default()),
        Box::new(storage::StorageSchemaDecode::default()),
        //Box::new(storage::JellyfishGetWithProof::default()),
        Box::new(storage::JellyfishGetWithProofWithDistinctLastNibble::default()),
        Box::new(storage::JellyfishGetRangeProof::default()),
        Box::new(storage::AccumulatorFrozenSubtreeHashes::default()),
        Box::new(storage::AccumulatorProof::default()),
        Box::new(storage::AccumulatorConsistencyProof::default()),
        Box::new(storage::AccumulatorRangeProof::default()),
        Box::new(storage::AccumulatorAppendMany::default()),
        Box::new(storage::AccumulatorAppendEmpty::default()),
        // Transaction
        Box::new(transaction::LanguageTransactionExecution::default()),
        Box::new(transaction::SignedTransactionTarget::default()),
        Box::new(transaction::MutatedSignedTransaction::default()),
        Box::new(transaction::TwoSignedTransactions::default()),
        // VM
        Box::new(vm::CompiledModuleTarget::default()),
    ];
    targets
        .into_iter()
        .map(|target| (target.name(), target))
        .collect()
});

impl FuzzTarget {
    /// The environment variable used for passing fuzz targets to child processes.
    pub(crate) const ENV_VAR: &'static str = "FUZZ_TARGET";

    /// Get the current fuzz target from the environment.
    pub fn from_env() -> Result<Self> {
        let name = env::var(Self::ENV_VAR)?;
        Self::by_name(&name).ok_or_else(|| format_err!("Unknown fuzz target '{}'", name))
    }

    /// Get a fuzz target by name.
    pub fn by_name(name: &str) -> Option<Self> {
        ALL_TARGETS.get(name).map(|target| FuzzTarget(&**target))
    }

    /// A list of all fuzz targets.
    pub fn all_targets() -> impl Iterator<Item = Self> {
        ALL_TARGETS.values().map(|target| FuzzTarget(&**target))
    }
}
<commit_msg>add the smt correctness fuzz target to the list<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{FuzzTarget, FuzzTargetImpl};
use anyhow::{format_err, Result};
use once_cell::sync::Lazy;
use std::{collections::BTreeMap, env};

// List fuzz target modules here.
mod consensus;
mod executor;
mod json_rpc_service;
mod mempool;
mod move_vm;
mod network;
mod proof;
mod safety_rules;
mod secure_storage_vault;
mod state_sync;
mod storage;
mod transaction;
mod vm;

static ALL_TARGETS: Lazy<BTreeMap<&'static str, Box<dyn FuzzTargetImpl>>> = Lazy::new(|| {
    // List fuzz targets here in this format:
    let targets: Vec<Box<dyn FuzzTargetImpl>> = vec![
        // Consensus
        Box::new(consensus::ConsensusProposal::default()),
        // Executor
        Box::new(executor::ExecuteAndCommitBlocks::default()),
        Box::new(executor::ExecuteAndCommitChunk::default()),
        // JSON RPC Service
        Box::new(json_rpc_service::JsonRpcSubmitTransactionRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetAccountRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetAccountTransactionRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetAccountTransactionsRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetTransactionsRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetEventsRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetMetadataRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetCurrenciesRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetStateProofRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetAccountStateWithProofRequest::default()),
        Box::new(json_rpc_service::JsonRpcGetNetworkStatusRequest::default()),
        // Mempool
        Box::new(mempool::MempoolIncomingTransactions::default()),
        // Move VM
        Box::new(move_vm::ValueTarget::default()),
        // Proof
        Box::new(proof::TestAccumulatorProofFuzzer::default()),
        Box::new(proof::SparseMerkleProofFuzzer::default()),
        Box::new(proof::TestAccumulatorRangeProofFuzzer::default()),
        Box::new(proof::TransactionInfoWithProofFuzzer::default()),
        Box::new(proof::AccountStateProofFuzzer::default()),
        Box::new(proof::EventProofFuzzer::default()),
        Box::new(proof::TransactionListProofFuzzer::default()),
        // Network
        Box::new(network::NetworkNoiseInitiator::default()),
        Box::new(network::NetworkNoiseResponder::default()),
        Box::new(network::NetworkNoiseStream::default()),
        Box::new(network::NetworkHandshakeExchange::default()),
        Box::new(network::NetworkHandshakeNegotiation::default()),
        Box::new(network::PeerNetworkMessagesReceive::default()),
        // Safety Rules Server (LSR)
        Box::new(safety_rules::SafetyRulesConstructAndSignVote::default()),
        Box::new(safety_rules::SafetyRulesInitialize::default()),
        Box::new(safety_rules::SafetyRulesHandleMessage::default()),
        Box::new(safety_rules::SafetyRulesSignProposal::default()),
        Box::new(safety_rules::SafetyRulesSignTimeout::default()),
        // Secure Storage Vault
        Box::new(secure_storage_vault::VaultGenericResponse::default()),
        Box::new(secure_storage_vault::VaultPolicyReadResponse::default()),
        Box::new(secure_storage_vault::VaultPolicyListResponse::default()),
        Box::new(secure_storage_vault::VaultSecretListResponse::default()),
        Box::new(secure_storage_vault::VaultSecretReadResponse::default()),
        Box::new(secure_storage_vault::VaultTokenCreateResponse::default()),
        Box::new(secure_storage_vault::VaultTokenRenewResponse::default()),
        Box::new(secure_storage_vault::VaultTransitCreateResponse::default()),
        Box::new(secure_storage_vault::VaultTransitExportResponse::default()),
        Box::new(secure_storage_vault::VaultTransitListResponse::default()),
        Box::new(secure_storage_vault::VaultTransitReadResponse::default()),
        Box::new(secure_storage_vault::VaultTransitRestoreResponse::default()),
        Box::new(secure_storage_vault::VaultTransitSignResponse::default()),
        Box::new(secure_storage_vault::VaultUnsealedResponse::default()),
        // State Sync
        Box::new(state_sync::StateSyncMsg::default()),
        // Storage
        // Box::new(storage::StorageSaveBlocks::default()),
        Box::new(storage::StorageSchemaDecode::default()),
        //Box::new(storage::JellyfishGetWithProof::default()),
        Box::new(storage::JellyfishGetWithProofWithDistinctLastNibble::default()),
        Box::new(storage::JellyfishGetRangeProof::default()),
        Box::new(storage::AccumulatorFrozenSubtreeHashes::default()),
        Box::new(storage::AccumulatorProof::default()),
        Box::new(storage::AccumulatorConsistencyProof::default()),
        Box::new(storage::AccumulatorRangeProof::default()),
        Box::new(storage::AccumulatorAppendMany::default()),
        Box::new(storage::AccumulatorAppendEmpty::default()),
        Box::new(storage::SparseMerkleCorrectness::default()),
        // Transaction
        Box::new(transaction::LanguageTransactionExecution::default()),
        Box::new(transaction::SignedTransactionTarget::default()),
        Box::new(transaction::MutatedSignedTransaction::default()),
        Box::new(transaction::TwoSignedTransactions::default()),
        // VM
        Box::new(vm::CompiledModuleTarget::default()),
    ];
    targets
        .into_iter()
        .map(|target| (target.name(), target))
        .collect()
});

impl FuzzTarget {
    /// The environment variable used for passing fuzz targets to child processes.
    pub(crate) const ENV_VAR: &'static str = "FUZZ_TARGET";

    /// Get the current fuzz target from the environment.
    pub fn from_env() -> Result<Self> {
        let name = env::var(Self::ENV_VAR)?;
        Self::by_name(&name).ok_or_else(|| format_err!("Unknown fuzz target '{}'", name))
    }

    /// Get a fuzz target by name.
    pub fn by_name(name: &str) -> Option<Self> {
        ALL_TARGETS.get(name).map(|target| FuzzTarget(&**target))
    }

    /// A list of all fuzz targets.
    pub fn all_targets() -> impl Iterator<Item = Self> {
        ALL_TARGETS.values().map(|target| FuzzTarget(&**target))
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

use super::Position;
use super::symbols::SymbolTable;

pub type Functions = HashMap<String, Node<Function>>;

#[derive(PartialEq, Debug)]
pub struct Node<I> {
    pub pos: Position,
    pub node: I,
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub enum Type {
    Void,
    Bool,
    Int,
    Float,
    Str,
}

#[derive(PartialEq, Debug)]
pub enum Literal {
    Bool(bool),
    Int(i32),
    Float(f32),
    Str(String),
}

impl Literal {
    pub fn get_type(&self) -> Type {
        match *self {
            Literal::Bool(_) => Type::Bool,
            Literal::Int(_) => Type::Int,
            Literal::Float(_) => Type::Float,
            Literal::Str(_) => Type::Str,
        }
    }
}

#[derive(Debug)]
pub struct Variable {
    pub id: u32,
    pub type_: Type,
    pub name: String,
}

impl Variable {
    pub fn stub(name: String) -> Variable {
        Variable {
            id: 0,
            type_: Type::Void,
            name: name,
        }
    }

    pub fn eq_name_type(&self, other: &Variable) -> bool {
        self.type_ == other.type_ && self.name == other.name
    }
}

impl PartialEq for Variable {
    fn eq(&self, other: &Variable) -> bool {
        if self.id == 0 || other.id == 0 {
            false
        } else {
            self.id == other.id
        }
    }
}

#[derive(PartialEq, Debug)]
pub enum UnaryOp {
    NOT,
    MINUS,
}

#[derive(PartialEq, Debug)]
pub enum BinaryOp {
    ADD,
    SUB,
    MUL,
    DIV,

    EQ,
    LT,
    LE,
}

#[derive(PartialEq, Debug)]
pub struct Program {
    pub functions: Functions,
}

#[derive(PartialEq, Debug)]
pub struct Function {
    pub name: String,
    pub filename: String,
    pub body: Node<Statement>,
    pub args: Vec<Rc<Variable>>,
    pub ret_type: Type,
    pub symbols: Rc<RefCell<SymbolTable>>,
}

#[derive(PartialEq, Debug)]
pub enum Statement {
    Expression {
        expr: Box<Node<Expression>>,
    },
    Declaration {
        var: Rc<Variable>,
    },
    Assignment {
        var: Rc<Variable>,
        expr: Box<Node<Expression>>,
    },
    If {
        cond: Box<Node<Expression>>,
        on_true: Box<Node<Statement>>,
        on_false: Option<Box<Node<Statement>>>,
    },
    While {
        cond: Box<Node<Expression>>,
        body: Box<Node<Statement>>,
    },
    Return {
        expr: Option<Box<Node<Expression>>>,
    },
    Compound {
        stmts: Vec<Node<Statement>>,
        symbols: Rc<RefCell<SymbolTable>>,
    },
}

impl Node<Statement> {
    pub fn visit_stmt(&self, do_stmt: &mut FnMut(&Node<Statement>) -> bool) -> bool {
        do_stmt(self) || return false;
        match self.node {
            Statement::If { ref on_true, ref on_false, .. } => {
                on_true.visit_stmt(do_stmt) ||
                if let Some(ref stmt) = *on_false {
                    stmt.visit_stmt(do_stmt)
                } else {
                    true
                }
            }
            Statement::While { ref body, .. } => body.visit_stmt(do_stmt),
            Statement::Compound { ref stmts, .. } => {
                for stmt in stmts.iter() {
                    stmt.visit_stmt(do_stmt) || return false;
                }
                true
            }
            _ => true,
        }
    }

    pub fn visit_expr(&self, do_expr: &mut FnMut(&Node<Expression>) -> bool) -> bool {
        self.visit_stmt(&mut |stmt| {
            match stmt.node {
                Statement::Expression { ref expr } => do_expr(expr),
                Statement::Assignment { ref expr, .. } => do_expr(expr),
                Statement::If { ref cond, .. } => do_expr(cond),
                Statement::While { ref cond, .. } => do_expr(cond),
                Statement::Return { expr: Some(ref expr) } => do_expr(expr),
                _ => true,
            }
        })
    }
}

#[derive(PartialEq, Debug)]
pub enum Expression {
    Literal {
        lit: Literal,
    },
    Variable {
        var: Rc<Variable>,
    },
    Call {
        function: String,
        args: Vec<Node<Expression>>,
    },
    Unary {
        op: UnaryOp,
        expr: Box<Node<Expression>>,
    },
    Binary {
        op: BinaryOp,
        left: Box<Node<Expression>>,
        right: Box<Node<Expression>>,
    },
    Parenthesis {
        expr: Box<Node<Expression>>,
    },
}

impl Node<Expression> {
    pub fn visit_expr(&self, do_expr: &mut FnMut(&Node<Expression>) -> bool) -> bool {
        do_expr(self) || return false;
        match self.node {
            Expression::Call { ref args, .. } => {
                for arg in args.iter() {
                    arg.visit_expr(do_expr) || return false;
                }
                true
            }
            Expression::Unary { ref expr, .. } => expr.visit_expr(do_expr),
            Expression::Parenthesis { ref expr } => expr.visit_expr(do_expr),
            Expression::Binary { ref left, ref right, .. } => {
                left.visit_expr(do_expr) || return false;
                right.visit_expr(do_expr)
            }
            _ => true,
        }
    }
}

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

    #[test]
    fn literals() {
        assert_eq!(Type::Bool, Literal::Bool(false).get_type());
        assert_eq!(Type::Int, Literal::Int(0).get_type());
        assert_eq!(Type::Float, Literal::Float(0.0).get_type());
    }
}
<commit_msg>Front: Add ast::Op<commit_after>use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

use super::Position;
use super::symbols::SymbolTable;

pub type Functions = HashMap<String, Node<Function>>;

#[derive(PartialEq, Debug)]
pub struct Node<I> {
    pub pos: Position,
    pub node: I,
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub enum Type {
    Void,
    Bool,
    Int,
    Float,
    Str,
}

#[derive(PartialEq, Debug)]
pub enum Literal {
    Bool(bool),
    Int(i32),
    Float(f32),
    Str(String),
}

impl Literal {
    pub fn get_type(&self) -> Type {
        match *self {
            Literal::Bool(_) => Type::Bool,
            Literal::Int(_) => Type::Int,
            Literal::Float(_) => Type::Float,
            Literal::Str(_) => Type::Str,
        }
    }
}

#[derive(Debug)]
pub struct Variable {
    pub id: u32,
    pub type_: Type,
    pub name: String,
}

impl Variable {
    pub fn stub(name: String) -> Variable {
        Variable {
            id: 0,
            type_: Type::Void,
            name: name,
        }
    }

    pub fn eq_name_type(&self, other: &Variable) -> bool {
        self.type_ == other.type_ && self.name == other.name
    }
}

impl PartialEq for Variable {
    fn eq(&self, other: &Variable) -> bool {
        if self.id == 0 || other.id == 0 {
            false
        } else {
            self.id == other.id
        }
    }
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub enum UnaryOp {
    NOT,
    MINUS,
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub enum BinaryOp {
    ADD,
    SUB,
    MUL,
    DIV,

    EQ,
    LT,
    LE,
}

#[derive(PartialEq, Debug, Clone, Copy)]
pub enum Op {
    UnaryOp(UnaryOp),
    BinaryOp(BinaryOp),
}

#[derive(PartialEq, Debug)]
pub struct Program {
    pub functions: Functions,
}

#[derive(PartialEq, Debug)]
pub struct Function {
    pub name: String,
    pub filename: String,
    pub body: Node<Statement>,
    pub args: Vec<Rc<Variable>>,
    pub ret_type: Type,
    pub symbols: Rc<RefCell<SymbolTable>>,
}

#[derive(PartialEq, Debug)]
pub enum Statement {
    Expression {
        expr: Box<Node<Expression>>,
    },
    Declaration {
        var: Rc<Variable>,
    },
    Assignment {
        var: Rc<Variable>,
        expr: Box<Node<Expression>>,
    },
    If {
        cond: Box<Node<Expression>>,
        on_true: Box<Node<Statement>>,
        on_false: Option<Box<Node<Statement>>>,
    },
    While {
        cond: Box<Node<Expression>>,
        body: Box<Node<Statement>>,
    },
    Return {
        expr: Option<Box<Node<Expression>>>,
    },
    Compound {
        stmts: Vec<Node<Statement>>,
        symbols: Rc<RefCell<SymbolTable>>,
    },
}

impl Node<Statement> {
    pub fn visit_stmt(&self, do_stmt: &mut FnMut(&Node<Statement>) -> bool) -> bool {
        do_stmt(self) || return false;
        match self.node {
            Statement::If { ref on_true, ref on_false, .. } => {
                on_true.visit_stmt(do_stmt) ||
                if let Some(ref stmt) = *on_false {
                    stmt.visit_stmt(do_stmt)
                } else {
                    true
                }
            }
            Statement::While { ref body, .. } => body.visit_stmt(do_stmt),
            Statement::Compound { ref stmts, .. } => {
                for stmt in stmts.iter() {
                    stmt.visit_stmt(do_stmt) || return false;
                }
                true
            }
            _ => true,
        }
    }

    pub fn visit_expr(&self, do_expr: &mut FnMut(&Node<Expression>) -> bool) -> bool {
        self.visit_stmt(&mut |stmt| {
            match stmt.node {
                Statement::Expression { ref expr } => do_expr(expr),
                Statement::Assignment { ref expr, .. } => do_expr(expr),
                Statement::If { ref cond, .. } => do_expr(cond),
                Statement::While { ref cond, .. } => do_expr(cond),
                Statement::Return { expr: Some(ref expr) } => do_expr(expr),
                _ => true,
            }
        })
    }
}

#[derive(PartialEq, Debug)]
pub enum Expression {
    Literal {
        lit: Literal,
    },
    Variable {
        var: Rc<Variable>,
    },
    Call {
        function: String,
        args: Vec<Node<Expression>>,
    },
    Unary {
        op: UnaryOp,
        expr: Box<Node<Expression>>,
    },
    Binary {
        op: BinaryOp,
        left: Box<Node<Expression>>,
        right: Box<Node<Expression>>,
    },
    Parenthesis {
        expr: Box<Node<Expression>>,
    },
}

impl Node<Expression> {
    pub fn visit_expr(&self, do_expr: &mut FnMut(&Node<Expression>) -> bool) -> bool {
        do_expr(self) || return false;
        match self.node {
            Expression::Call { ref args, .. } => {
                for arg in args.iter() {
                    arg.visit_expr(do_expr) || return false;
                }
                true
            }
            Expression::Unary { ref expr, .. } => expr.visit_expr(do_expr),
            Expression::Parenthesis { ref expr } => expr.visit_expr(do_expr),
            Expression::Binary { ref left, ref right, .. } => {
                left.visit_expr(do_expr) || return false;
                right.visit_expr(do_expr)
            }
            _ => true,
        }
    }
}

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

    #[test]
    fn literals() {
        assert_eq!(Type::Bool, Literal::Bool(false).get_type());
        assert_eq!(Type::Int, Literal::Int(0).get_type());
        assert_eq!(Type::Float, Literal::Float(0.0).get_type());
    }
}
<|endoftext|>


<commit_before>use std::{cmp::min, env, fs::File, io, mem, time::Duration};

use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, rngs::SmallRng, Rng, SeedableRng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 50;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    filename: String,
    client: reqwest::blocking::Client,
    rng: SmallRng,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(filename: String) -> Importer {
        Importer {
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            rng: SmallRng::from_seed([
                0x19, 0x29, 0xab, 0x17, 0xc6, 0xfa, 0xb0, 0xe9, 0x4b, 0x44, 0xd8, 0x07, 0x09, 0xbf,
                0x1d, 0x87, 0xbd, 0xd8, 0xb3, 0x2f, 0xe1, 0xe2, 0xa0, 0x1a, 0x9e, 0x30, 0x98, 0xd7,
                0xef, 0xd5, 0x7a, 0x1d,
            ]),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put("http://127.0.0.1:9001/import/lichess")
            .json(&self.batch)
            .send()
            .expect("send batch");

        println!(
            "{}: {}: {} - {}",
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(false, |name| name != "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.0,

                _ if rating >= 2500 => 1.0,

                Speed::Rapid if rating >= 2200 => 1.0,
                Speed::Rapid if rating >= 2000 => 0.30,
                Speed::Rapid if rating >= 1800 => 0.18,
                Speed::Rapid if rating >= 1600 => 0.16,

                Speed::Blitz if rating >= 2200 => 0.16,
                Speed::Blitz if rating >= 2000 => 0.07,
                Speed::Blitz if rating >= 1800 => 0.05,
                Speed::Blitz if rating >= 1600 => 0.05,

                Speed::Bullet if rating >= 2200 => 0.16,
                Speed::Bullet if rating >= 2000 => 0.07,
                Speed::Bullet if rating >= 1800 => 0.05,
                Speed::Bullet if rating >= 1600 => 0.05,

                Speed::UltraBullet => 1.0,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.0
            } else {
                0.5
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= self.rng.sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

fn main() -> Result<(), io::Error> {
    for arg in env::args().skip(1) {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.ends_with(".bz2") {
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(arg);
        reader.read_all(&mut importer)?;
    }

    Ok(())
}
<commit_msg>fix bullet sampling rates<commit_after>use std::{cmp::min, env, fs::File, io, mem, time::Duration};

use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, rngs::SmallRng, Rng, SeedableRng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 50;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    filename: String,
    client: reqwest::blocking::Client,
    rng: SmallRng,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(filename: String) -> Importer {
        Importer {
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            rng: SmallRng::from_seed([
                0x19, 0x29, 0xab, 0x17, 0xc6, 0xfa, 0xb0, 0xe9, 0x4b, 0x44, 0xd8, 0x07, 0x09, 0xbf,
                0x1d, 0x87, 0xbd, 0xd8, 0xb3, 0x2f, 0xe1, 0xe2, 0xa0, 0x1a, 0x9e, 0x30, 0x98, 0xd7,
                0xef, 0xd5, 0x7a, 0x1d,
            ]),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put("http://127.0.0.1:9001/import/lichess")
            .json(&self.batch)
            .send()
            .expect("send batch");

        println!(
            "{}: {}: {} - {}",
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(false, |name| name != "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.0,

                _ if rating >= 2500 => 1.0,

                Speed::Rapid if rating >= 2200 => 1.0,
                Speed::Rapid if rating >= 2000 => 0.30,
                Speed::Rapid if rating >= 1800 => 0.18,
                Speed::Rapid if rating >= 1600 => 0.16,

                Speed::Blitz if rating >= 2200 => 0.16,
                Speed::Blitz if rating >= 2000 => 0.07,
                Speed::Blitz if rating >= 1800 => 0.05,
                Speed::Blitz if rating >= 1600 => 0.05,

                Speed::Bullet if rating >= 2200 => 0.20,
                Speed::Bullet if rating >= 2000 => 0.11,
                Speed::Bullet if rating >= 1800 => 0.07,
                Speed::Bullet if rating >= 1600 => 0.07,

                Speed::UltraBullet => 1.0,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.0
            } else {
                0.5
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= self.rng.sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

fn main() -> Result<(), io::Error> {
    for arg in env::args().skip(1) {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.ends_with(".bz2") {
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(arg);
        reader.read_all(&mut importer)?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>extern crate rand;

use std::cmp::Ordering;

#[derive(Debug, Clone)]
pub struct Gene {
    in_neuron_id: usize,
    out_neuron_id: usize,
    weight: f64,
    enabled: bool,
    innovation: u64,
}

impl Eq for Gene {}

impl PartialEq for Gene {
    fn eq(&self, other: &Gene) -> bool {
        self.in_neuron_id == other.in_neuron_id && self.out_neuron_id == other.out_neuron_id
    }
}

impl Ord for Gene {
    fn cmp(&self, other: &Gene) -> Ordering {
        if self == other {
            Ordering::Equal
        } else if self.in_neuron_id == other.in_neuron_id {
            if self.out_neuron_id > other.out_neuron_id {
                Ordering::Greater
            } else {
                Ordering::Less
            }
        } else if self.in_neuron_id > other.in_neuron_id {
            Ordering::Greater
        } else {
            Ordering::Less
        }
    }
}

impl PartialOrd for Gene {
    fn partial_cmp(&self, other: &Gene) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Gene {
    pub fn generate_weight() -> f64 {
        rand::random::<f64>() * 2f64 - 1f64
    }
}

impl Default for Gene {
    fn default() -> Gene {
        Gene {
            in_neuron_id: 1,
            out_neuron_id: 1,
            weight: Gene::generate_weight(),
            enabled: true,
        }
    }
}

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

    fn g(n_in: usize, n_out: usize) -> Gene {
        Gene {
            in_neuron_id: n_in,
            out_neuron_id: n_out,
            ..Gene::default()
        }
    }

    #[test]
    fn should_be_able_to_binary_search_for_a_gene() {
        let mut genome = vec![g(0, 1), g(0, 2), g(3, 2), g(2, 3), g(1, 5)];
        genome.sort();
        genome.binary_search(&g(0, 1)).unwrap();
        genome.binary_search(&g(0, 2)).unwrap();
        genome.binary_search(&g(1, 5)).unwrap();
        genome.binary_search(&g(2, 3)).unwrap();
        genome.binary_search(&g(3, 2)).unwrap();
    }
}
<commit_msg>fix/gene Redo earlier edit. This PR should be rustfmt only<commit_after>extern crate rand;

use std::cmp::Ordering;

#[derive(Debug, Clone)]
pub struct Gene {
    pub in_neuron_id: usize,
    pub out_neuron_id: usize,
    pub weight: f64,
    pub enabled: bool,
}

impl Eq for Gene {}

impl PartialEq for Gene {
    fn eq(&self, other: &Gene) -> bool {
        self.in_neuron_id == other.in_neuron_id && self.out_neuron_id == other.out_neuron_id
    }
}

impl Ord for Gene {
    fn cmp(&self, other: &Gene) -> Ordering {
        if self == other {
            Ordering::Equal
        } else if self.in_neuron_id == other.in_neuron_id {
            if self.out_neuron_id > other.out_neuron_id {
                Ordering::Greater
            } else {
                Ordering::Less
            }
        } else if self.in_neuron_id > other.in_neuron_id {
            Ordering::Greater
        } else {
            Ordering::Less
        }
    }
}

impl PartialOrd for Gene {
    fn partial_cmp(&self, other: &Gene) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Gene {
    pub fn generate_weight() -> f64 {
        rand::random::<f64>() * 2f64 - 1f64
    }
}

impl Default for Gene {
    fn default() -> Gene {
        Gene {
            in_neuron_id: 1,
            out_neuron_id: 1,
            weight: Gene::generate_weight(),
            enabled: true,
        }
    }
}

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

    fn g(n_in: usize, n_out: usize) -> Gene {
        Gene {
            in_neuron_id: n_in,
            out_neuron_id: n_out,
            ..Gene::default()
        }
    }

    #[test]
    fn should_be_able_to_binary_search_for_a_gene() {
        let mut genome = vec![g(0, 1), g(0, 2), g(3, 2), g(2, 3), g(1, 5)];
        genome.sort();
        genome.binary_search(&g(0, 1)).unwrap();
        genome.binary_search(&g(0, 2)).unwrap();
        genome.binary_search(&g(1, 5)).unwrap();
        genome.binary_search(&g(2, 3)).unwrap();
        genome.binary_search(&g(3, 2)).unwrap();
    }
}
<|endoftext|>


<commit_before>use types::{
    ValidatorRecord,
    ValidatorStatus,
};

use super::proof_of_possession::verify_proof_of_possession;
use super::registration::ValidatorRegistration;

/// The size of a validators deposit in GWei.
pub const DEPOSIT_GWEI: u64 = 32_000_000_000;

/// Inducts validators into a `CrystallizedState`.
pub struct ValidatorInductor<'a> {
    pub current_slot: u64,
    pub shard_count: u16,
    validators: &'a mut Vec<ValidatorRecord>,
    empty_validator_start: usize,
}

#[derive(Debug, PartialEq, Clone)]
pub enum ValidatorInductionError {
    InvalidShard,
    InvaidProofOfPossession,
}

impl<'a> ValidatorInductor<'a> {
    /// Attempt to induct a validator into the CrystallizedState.
    ///
    /// Returns an error if the registration is invalid, otherwise returns the index of the
    /// validator in `CrystallizedState.validators`.
    pub fn induct(&mut self, rego: &ValidatorRegistration)
        -> Result<usize, ValidatorInductionError>
    {
        let v = self.process_registration(rego)?;
        Ok(self.add_validator(v))
    }

    /// Verify a `ValidatorRegistration` and return a `ValidatorRecord` if valid.
    fn process_registration(&self, r: &ValidatorRegistration)
        -> Result<ValidatorRecord, ValidatorInductionError>
    {
        /*
         * Ensure withdrawal shard is not too high.
         */
        if r.withdrawal_shard > self.shard_count {
            return Err(ValidatorInductionError::InvalidShard)
        }

        /*
         * Prove validator has knowledge of their secret key.
         */
        if !verify_proof_of_possession(&r.proof_of_possession, &r.pubkey) {
            return Err(ValidatorInductionError::InvaidProofOfPossession)
        }

        Ok(ValidatorRecord {
            pubkey: r.pubkey.clone(),
            withdrawal_shard: r.withdrawal_shard,
            withdrawal_address: r.withdrawal_address,
            randao_commitment: r.randao_commitment,
            randao_last_change: self.current_slot,
            balance: DEPOSIT_GWEI,
            status: ValidatorStatus::PendingActivation as u8,
            exit_slot: 0,
        })
    }

    /// Returns the index of the first `ValidatorRecord` in the `CrystallizedState` where
    /// `validator.status == Withdrawn`. If no such record exists, `None` is returned.
    fn first_withdrawn_validator(&mut self)
        -> Option<usize>
    {
        for i in self.empty_validator_start..self.validators.len() {
            if self.validators[i].status == ValidatorStatus::Withdrawn as u8 {
                self.empty_validator_start = i + 1;
                return Some(i)
            }
        }
        None
    }

    /// Adds a `ValidatorRecord` to the `CrystallizedState` by replacing first validator where
    /// `validator.status == Withdraw`. If no such withdrawn validator exists, adds the new
    /// validator to the end of the list.
    fn add_validator(&mut self, v: ValidatorRecord)
        -> usize
    {
        match self.first_withdrawn_validator() {
            Some(i) => {
                self.validators[i] = v;
                i
            }
            None => {
                self.validators.push(v);
                self.validators.len() - 1
            }
        }
    }
}

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

    use bls::{
        Keypair,
        Signature,
    };
    use types::{
        Address,
        Hash256,
    };
    use hashing::proof_of_possession_hash;

    /// Generate a proof of possession for some keypair.
    fn get_proof_of_possession(kp: &Keypair) -> Signature {
        let pop_message = proof_of_possession_hash(&kp.pk.as_bytes());
        Signature::new_hashed(&pop_message, &kp.sk)
    }

    /// Generate a basic working ValidatorRegistration for use in tests.
    fn get_registration() -> ValidatorRegistration {
        let kp = Keypair::random();
        ValidatorRegistration {
            pubkey: kp.pk.clone(),
            withdrawal_shard: 0,
            withdrawal_address: Address::zero(),
            randao_commitment: Hash256::zero(),
            proof_of_possession: get_proof_of_possession(&kp),
        }
    }

    /// Induct a validator using the ValidatorInductor, return the result.
    fn do_induction(validator_rego: &ValidatorRegistration,
                    validators: &mut Vec<ValidatorRecord>,
                    current_slot: u64,
                    shard_count: u16)
        -> Result<usize, ValidatorInductionError>
    {
        let mut inductor = ValidatorInductor {
            current_slot,
            shard_count,
            empty_validator_start: 0,
            validators,
        };
        inductor.induct(&validator_rego)
    }

    #[test]
    fn test_validator_inductor_valid_empty_validators() {
        let mut validators = vec![];

        let r = get_registration();

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result.unwrap(), 0);
        assert_eq!(r, validators[0]);
        assert_eq!(validators.len(), 1);
    }

    #[test]
    fn test_validator_inductor_valid_all_active_validators() {
        let mut validators = vec![];
        for _ in 0..5 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Active as u8;
            validators.push(v);
        }

        let r = get_registration();

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result.unwrap(), 5);
        assert_eq!(r, validators[5]);
        assert_eq!(validators.len(), 6);
    }

    #[test]
    fn test_validator_inductor_valid_all_second_validator_withdrawn() {
        let mut validators = vec![];
        let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
        v.status = ValidatorStatus::Active as u8;
        validators.push(v);
        for _ in 0..4 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Withdrawn as u8;
            validators.push(v);
        }

        let r = get_registration();

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result.unwrap(), 1);
        assert_eq!(r, validators[1]);
        assert_eq!(validators.len(), 5);
    }

    #[test]
    fn test_validator_inductor_valid_all_withdrawn_validators() {
        let mut validators = vec![];
        for _ in 0..5 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Withdrawn as u8;
            validators.push(v);
        }

        /*
         * Ensure the first validator gets the 0'th slot
         */
        let r = get_registration();
        let result = do_induction(&r, &mut validators, 0, 1024);
        assert_eq!(result.unwrap(), 0);
        assert_eq!(r, validators[0]);
        assert_eq!(validators.len(), 5);

        /*
         * Ensure the second validator gets the 1'st slot
         */
        let r_two = get_registration();
        let result = do_induction(&r_two, &mut validators, 0, 1024);
        assert_eq!(result.unwrap(), 1);
        assert_eq!(r_two, validators[1]);
        assert_eq!(validators.len(), 5);
    }

    #[test]
    fn test_validator_inductor_shard_too_high() {
        let mut validators = vec![];

        let mut r = get_registration();
        r.withdrawal_shard = 1025;

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result, Err(ValidatorInductionError::InvalidShard));
        assert_eq!(validators.len(), 0);
    }

    #[test]
    fn test_validator_inductor_shard_proof_of_possession_failure() {
        let mut validators = vec![];

        let mut r = get_registration();
        let kp = Keypair::random();
        r.proof_of_possession = get_proof_of_possession(&kp);

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result, Err(ValidatorInductionError::InvaidProofOfPossession));
        assert_eq!(validators.len(), 0);
    }
}
<commit_msg>Add new method for ValidatorInductor<commit_after>use types::{
    ValidatorRecord,
    ValidatorStatus,
};

use super::proof_of_possession::verify_proof_of_possession;
use super::registration::ValidatorRegistration;

/// The size of a validators deposit in GWei.
pub const DEPOSIT_GWEI: u64 = 32_000_000_000;

/// Inducts validators into a `CrystallizedState`.
pub struct ValidatorInductor<'a> {
    pub current_slot: u64,
    pub shard_count: u16,
    validators: &'a mut Vec<ValidatorRecord>,
    empty_validator_start: usize,
}

#[derive(Debug, PartialEq, Clone)]
pub enum ValidatorInductionError {
    InvalidShard,
    InvaidProofOfPossession,
}

impl<'a> ValidatorInductor<'a> {
    pub fn new(current_slot: u64, shard_count: u16, validators: &'a mut Vec<ValidatorRecord>)
        -> Self
    {
        Self {
            current_slot,
            shard_count,
            validators,
            empty_validator_start: 0,
        }
    }

    /// Attempt to induct a validator into the CrystallizedState.
    ///
    /// Returns an error if the registration is invalid, otherwise returns the index of the
    /// validator in `CrystallizedState.validators`.
    pub fn induct(&mut self, rego: &ValidatorRegistration)
        -> Result<usize, ValidatorInductionError>
    {
        let v = self.process_registration(rego)?;
        Ok(self.add_validator(v))
    }

    /// Verify a `ValidatorRegistration` and return a `ValidatorRecord` if valid.
    fn process_registration(&self, r: &ValidatorRegistration)
        -> Result<ValidatorRecord, ValidatorInductionError>
    {
        /*
         * Ensure withdrawal shard is not too high.
         */
        if r.withdrawal_shard > self.shard_count {
            return Err(ValidatorInductionError::InvalidShard)
        }

        /*
         * Prove validator has knowledge of their secret key.
         */
        if !verify_proof_of_possession(&r.proof_of_possession, &r.pubkey) {
            return Err(ValidatorInductionError::InvaidProofOfPossession)
        }

        Ok(ValidatorRecord {
            pubkey: r.pubkey.clone(),
            withdrawal_shard: r.withdrawal_shard,
            withdrawal_address: r.withdrawal_address,
            randao_commitment: r.randao_commitment,
            randao_last_change: self.current_slot,
            balance: DEPOSIT_GWEI,
            status: ValidatorStatus::PendingActivation as u8,
            exit_slot: 0,
        })
    }

    /// Returns the index of the first `ValidatorRecord` in the `CrystallizedState` where
    /// `validator.status == Withdrawn`. If no such record exists, `None` is returned.
    fn first_withdrawn_validator(&mut self)
        -> Option<usize>
    {
        for i in self.empty_validator_start..self.validators.len() {
            if self.validators[i].status == ValidatorStatus::Withdrawn as u8 {
                self.empty_validator_start = i + 1;
                return Some(i)
            }
        }
        None
    }

    /// Adds a `ValidatorRecord` to the `CrystallizedState` by replacing first validator where
    /// `validator.status == Withdraw`. If no such withdrawn validator exists, adds the new
    /// validator to the end of the list.
    fn add_validator(&mut self, v: ValidatorRecord)
        -> usize
    {
        match self.first_withdrawn_validator() {
            Some(i) => {
                self.validators[i] = v;
                i
            }
            None => {
                self.validators.push(v);
                self.validators.len() - 1
            }
        }
    }
}

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

    use bls::{
        Keypair,
        Signature,
    };
    use types::{
        Address,
        Hash256,
    };
    use hashing::proof_of_possession_hash;

    /// Generate a proof of possession for some keypair.
    fn get_proof_of_possession(kp: &Keypair) -> Signature {
        let pop_message = proof_of_possession_hash(&kp.pk.as_bytes());
        Signature::new_hashed(&pop_message, &kp.sk)
    }

    /// Generate a basic working ValidatorRegistration for use in tests.
    fn get_registration() -> ValidatorRegistration {
        let kp = Keypair::random();
        ValidatorRegistration {
            pubkey: kp.pk.clone(),
            withdrawal_shard: 0,
            withdrawal_address: Address::zero(),
            randao_commitment: Hash256::zero(),
            proof_of_possession: get_proof_of_possession(&kp),
        }
    }

    /// Induct a validator using the ValidatorInductor, return the result.
    fn do_induction(validator_rego: &ValidatorRegistration,
                    validators: &mut Vec<ValidatorRecord>,
                    current_slot: u64,
                    shard_count: u16)
        -> Result<usize, ValidatorInductionError>
    {
        let mut inductor = ValidatorInductor::new(
            current_slot,
            shard_count,
            validators);
        inductor.induct(&validator_rego)
    }

    #[test]
    fn test_validator_inductor_valid_empty_validators() {
        let mut validators = vec![];

        let r = get_registration();

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result.unwrap(), 0);
        assert_eq!(r, validators[0]);
        assert_eq!(validators.len(), 1);
    }

    #[test]
    fn test_validator_inductor_valid_all_active_validators() {
        let mut validators = vec![];
        for _ in 0..5 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Active as u8;
            validators.push(v);
        }

        let r = get_registration();

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result.unwrap(), 5);
        assert_eq!(r, validators[5]);
        assert_eq!(validators.len(), 6);
    }

    #[test]
    fn test_validator_inductor_valid_all_second_validator_withdrawn() {
        let mut validators = vec![];
        let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
        v.status = ValidatorStatus::Active as u8;
        validators.push(v);
        for _ in 0..4 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Withdrawn as u8;
            validators.push(v);
        }

        let r = get_registration();

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result.unwrap(), 1);
        assert_eq!(r, validators[1]);
        assert_eq!(validators.len(), 5);
    }

    #[test]
    fn test_validator_inductor_valid_all_withdrawn_validators() {
        let mut validators = vec![];
        for _ in 0..5 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Withdrawn as u8;
            validators.push(v);
        }

        /*
         * Ensure the first validator gets the 0'th slot
         */
        let r = get_registration();
        let result = do_induction(&r, &mut validators, 0, 1024);
        assert_eq!(result.unwrap(), 0);
        assert_eq!(r, validators[0]);
        assert_eq!(validators.len(), 5);

        /*
         * Ensure the second validator gets the 1'st slot
         */
        let r_two = get_registration();
        let result = do_induction(&r_two, &mut validators, 0, 1024);
        assert_eq!(result.unwrap(), 1);
        assert_eq!(r_two, validators[1]);
        assert_eq!(validators.len(), 5);
    }

    #[test]
    fn test_validator_inductor_shard_too_high() {
        let mut validators = vec![];

        let mut r = get_registration();
        r.withdrawal_shard = 1025;

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result, Err(ValidatorInductionError::InvalidShard));
        assert_eq!(validators.len(), 0);
    }

    #[test]
    fn test_validator_inductor_shard_proof_of_possession_failure() {
        let mut validators = vec![];

        let mut r = get_registration();
        let kp = Keypair::random();
        r.proof_of_possession = get_proof_of_possession(&kp);

        let result = do_induction(&r, &mut validators, 0, 1024);

        assert_eq!(result, Err(ValidatorInductionError::InvaidProofOfPossession));
        assert_eq!(validators.len(), 0);
    }
}
<|endoftext|>


<commit_before>use bls::verify_proof_of_possession;
use types::{ValidatorRecord, ValidatorRegistration, ValidatorStatus};

/// The size of a validators deposit in GWei.
pub const DEPOSIT_GWEI: u64 = 32_000_000_000;

/// Inducts validators into a `CrystallizedState`.
pub struct ValidatorInductor {
    pub current_slot: u64,
    pub shard_count: u16,
    validators: Vec<ValidatorRecord>,
    empty_validator_start: usize,
}

#[derive(Debug, PartialEq, Clone)]
pub enum ValidatorInductionError {
    InvalidShard,
    InvaidProofOfPossession,
}

impl ValidatorInductor {
    pub fn new(current_slot: u64, shard_count: u16, validators: Vec<ValidatorRecord>) -> Self {
        Self {
            current_slot,
            shard_count,
            validators,
            empty_validator_start: 0,
        }
    }

    /// Attempt to induct a validator into the CrystallizedState.
    ///
    /// Returns an error if the registration is invalid, otherwise returns the index of the
    /// validator in `CrystallizedState.validators`.
    pub fn induct(
        &mut self,
        rego: &ValidatorRegistration,
        status: ValidatorStatus,
    ) -> Result<usize, ValidatorInductionError> {
        let v = self.process_registration(rego, status)?;
        Ok(self.add_validator(v))
    }

    /// Verify a `ValidatorRegistration` and return a `ValidatorRecord` if valid.
    fn process_registration(
        &self,
        r: &ValidatorRegistration,
        status: ValidatorStatus,
    ) -> Result<ValidatorRecord, ValidatorInductionError> {
        /*
         * Ensure withdrawal shard is not too high.
         */
        if r.withdrawal_shard > self.shard_count {
            return Err(ValidatorInductionError::InvalidShard);
        }

        /*
         * Prove validator has knowledge of their secret key.
         */
        if !verify_proof_of_possession(&r.proof_of_possession, &r.pubkey) {
            return Err(ValidatorInductionError::InvaidProofOfPossession);
        }

        Ok(ValidatorRecord {
            pubkey: r.pubkey.clone(),
            withdrawal_shard: r.withdrawal_shard,
            withdrawal_address: r.withdrawal_address,
            randao_commitment: r.randao_commitment,
            randao_last_change: self.current_slot,
            balance: DEPOSIT_GWEI,
            status: status as u8,
            exit_slot: 0,
        })
    }

    /// Returns the index of the first `ValidatorRecord` in the `CrystallizedState` where
    /// `validator.status == Withdrawn`. If no such record exists, `None` is returned.
    fn first_withdrawn_validator(&mut self) -> Option<usize> {
        for i in self.empty_validator_start..self.validators.len() {
            if self.validators[i].status == ValidatorStatus::Withdrawn as u8 {
                self.empty_validator_start = i + 1;
                return Some(i);
            }
        }
        None
    }

    /// Adds a `ValidatorRecord` to the `CrystallizedState` by replacing first validator where
    /// `validator.status == Withdraw`. If no such withdrawn validator exists, adds the new
    /// validator to the end of the list.
    fn add_validator(&mut self, v: ValidatorRecord) -> usize {
        match self.first_withdrawn_validator() {
            Some(i) => {
                self.validators[i] = v;
                i
            }
            None => {
                self.validators.push(v);
                self.validators.len() - 1
            }
        }
    }

    pub fn to_vec(self) -> Vec<ValidatorRecord> {
        self.validators
    }
}

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

    use bls::{Keypair, Signature};
    use hashing::proof_of_possession_hash;
    use types::{Address, Hash256};

    fn registration_equals_record(reg: &ValidatorRegistration, rec: &ValidatorRecord) -> bool {
        (reg.pubkey == rec.pubkey)
            & (reg.withdrawal_shard == rec.withdrawal_shard)
            & (reg.withdrawal_address == rec.withdrawal_address)
            & (reg.randao_commitment == rec.randao_commitment)
            & (verify_proof_of_possession(&reg.proof_of_possession, &rec.pubkey))
    }

    /// Generate a proof of possession for some keypair.
    fn get_proof_of_possession(kp: &Keypair) -> Signature {
        let pop_message = proof_of_possession_hash(&kp.pk.as_bytes());
        Signature::new_hashed(&pop_message, &kp.sk)
    }

    /// Generate a basic working ValidatorRegistration for use in tests.
    fn get_registration() -> ValidatorRegistration {
        let kp = Keypair::random();
        ValidatorRegistration {
            pubkey: kp.pk.clone(),
            withdrawal_shard: 0,
            withdrawal_address: Address::zero(),
            randao_commitment: Hash256::zero(),
            proof_of_possession: get_proof_of_possession(&kp),
        }
    }

    #[test]
    fn test_validator_inductor_valid_empty_validators() {
        let validators = vec![];

        let r = get_registration();

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(result.unwrap(), 0);
        assert!(registration_equals_record(&r, &validators[0]));
        assert_eq!(validators.len(), 1);
    }

    #[test]
    fn test_validator_inductor_status() {
        let validators = vec![];

        let r = get_registration();

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let _ = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let _ = inductor.induct(&r, ValidatorStatus::Active);
        let validators = inductor.to_vec();

        assert!(validators[0].status == ValidatorStatus::PendingActivation as u8);
        assert!(validators[1].status == ValidatorStatus::Active as u8);
        assert_eq!(validators.len(), 2);
    }

    #[test]
    fn test_validator_inductor_valid_all_active_validators() {
        let mut validators = vec![];
        for _ in 0..5 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Active as u8;
            validators.push(v);
        }

        let r = get_registration();

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(result.unwrap(), 5);
        assert!(registration_equals_record(&r, &validators[5]));
        assert_eq!(validators.len(), 6);
    }

    #[test]
    fn test_validator_inductor_valid_all_second_validator_withdrawn() {
        let mut validators = vec![];
        let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
        v.status = ValidatorStatus::Active as u8;
        validators.push(v);
        for _ in 0..4 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Withdrawn as u8;
            validators.push(v);
        }

        let r = get_registration();

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(result.unwrap(), 1);
        assert!(registration_equals_record(&r, &validators[1]));
        assert_eq!(validators.len(), 5);
    }

    #[test]
    fn test_validator_inductor_valid_all_withdrawn_validators() {
        let mut validators = vec![];
        for _ in 0..5 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Withdrawn as u8;
            validators.push(v);
        }

        /*
         * Ensure the first validator gets the 0'th slot
         */
        let r = get_registration();
        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();
        assert_eq!(result.unwrap(), 0);
        assert!(registration_equals_record(&r, &validators[0]));

        /*
         * Ensure the second validator gets the 1'st slot
         */
        let r_two = get_registration();
        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r_two, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();
        assert_eq!(result.unwrap(), 1);
        assert!(registration_equals_record(&r_two, &validators[1]));
        assert_eq!(validators.len(), 5);
    }

    #[test]
    fn test_validator_inductor_shard_too_high() {
        let validators = vec![];

        let mut r = get_registration();
        r.withdrawal_shard = 1025;

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(result, Err(ValidatorInductionError::InvalidShard));
        assert_eq!(validators.len(), 0);
    }

    #[test]
    fn test_validator_inductor_shard_proof_of_possession_failure() {
        let validators = vec![];

        let mut r = get_registration();
        let kp = Keypair::random();
        r.proof_of_possession = get_proof_of_possession(&kp);

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(
            result,
            Err(ValidatorInductionError::InvaidProofOfPossession)
        );
        assert_eq!(validators.len(), 0);
    }
}
<commit_msg>update function we removed so test compiles<commit_after>use bls::verify_proof_of_possession;
use types::{ValidatorRecord, ValidatorRegistration, ValidatorStatus};

/// The size of a validators deposit in GWei.
pub const DEPOSIT_GWEI: u64 = 32_000_000_000;

/// Inducts validators into a `CrystallizedState`.
pub struct ValidatorInductor {
    pub current_slot: u64,
    pub shard_count: u16,
    validators: Vec<ValidatorRecord>,
    empty_validator_start: usize,
}

#[derive(Debug, PartialEq, Clone)]
pub enum ValidatorInductionError {
    InvalidShard,
    InvaidProofOfPossession,
}

impl ValidatorInductor {
    pub fn new(current_slot: u64, shard_count: u16, validators: Vec<ValidatorRecord>) -> Self {
        Self {
            current_slot,
            shard_count,
            validators,
            empty_validator_start: 0,
        }
    }

    /// Attempt to induct a validator into the CrystallizedState.
    ///
    /// Returns an error if the registration is invalid, otherwise returns the index of the
    /// validator in `CrystallizedState.validators`.
    pub fn induct(
        &mut self,
        rego: &ValidatorRegistration,
        status: ValidatorStatus,
    ) -> Result<usize, ValidatorInductionError> {
        let v = self.process_registration(rego, status)?;
        Ok(self.add_validator(v))
    }

    /// Verify a `ValidatorRegistration` and return a `ValidatorRecord` if valid.
    fn process_registration(
        &self,
        r: &ValidatorRegistration,
        status: ValidatorStatus,
    ) -> Result<ValidatorRecord, ValidatorInductionError> {
        /*
         * Ensure withdrawal shard is not too high.
         */
        if r.withdrawal_shard > self.shard_count {
            return Err(ValidatorInductionError::InvalidShard);
        }

        /*
         * Prove validator has knowledge of their secret key.
         */
        if !verify_proof_of_possession(&r.proof_of_possession, &r.pubkey) {
            return Err(ValidatorInductionError::InvaidProofOfPossession);
        }

        Ok(ValidatorRecord {
            pubkey: r.pubkey.clone(),
            withdrawal_shard: r.withdrawal_shard,
            withdrawal_address: r.withdrawal_address,
            randao_commitment: r.randao_commitment,
            randao_last_change: self.current_slot,
            balance: DEPOSIT_GWEI,
            status: status as u8,
            exit_slot: 0,
        })
    }

    /// Returns the index of the first `ValidatorRecord` in the `CrystallizedState` where
    /// `validator.status == Withdrawn`. If no such record exists, `None` is returned.
    fn first_withdrawn_validator(&mut self) -> Option<usize> {
        for i in self.empty_validator_start..self.validators.len() {
            if self.validators[i].status == ValidatorStatus::Withdrawn as u8 {
                self.empty_validator_start = i + 1;
                return Some(i);
            }
        }
        None
    }

    /// Adds a `ValidatorRecord` to the `CrystallizedState` by replacing first validator where
    /// `validator.status == Withdraw`. If no such withdrawn validator exists, adds the new
    /// validator to the end of the list.
    fn add_validator(&mut self, v: ValidatorRecord) -> usize {
        match self.first_withdrawn_validator() {
            Some(i) => {
                self.validators[i] = v;
                i
            }
            None => {
                self.validators.push(v);
                self.validators.len() - 1
            }
        }
    }

    pub fn to_vec(self) -> Vec<ValidatorRecord> {
        self.validators
    }
}

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

    use bls::{Keypair, Signature};
    use hashing::canonical_hash;
    use types::{Address, Hash256};

    fn registration_equals_record(reg: &ValidatorRegistration, rec: &ValidatorRecord) -> bool {
        (reg.pubkey == rec.pubkey)
            & (reg.withdrawal_shard == rec.withdrawal_shard)
            & (reg.withdrawal_address == rec.withdrawal_address)
            & (reg.randao_commitment == rec.randao_commitment)
            & (verify_proof_of_possession(&reg.proof_of_possession, &rec.pubkey))
    }

    /// Generate a proof of possession for some keypair.
    fn get_proof_of_possession(kp: &Keypair) -> Signature {
        let pop_message = canonical_hash(&kp.pk.as_bytes());
        Signature::new_hashed(&pop_message, &kp.sk)
    }

    /// Generate a basic working ValidatorRegistration for use in tests.
    fn get_registration() -> ValidatorRegistration {
        let kp = Keypair::random();
        ValidatorRegistration {
            pubkey: kp.pk.clone(),
            withdrawal_shard: 0,
            withdrawal_address: Address::zero(),
            randao_commitment: Hash256::zero(),
            proof_of_possession: get_proof_of_possession(&kp),
        }
    }

    #[test]
    fn test_validator_inductor_valid_empty_validators() {
        let validators = vec![];

        let r = get_registration();

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(result.unwrap(), 0);
        assert!(registration_equals_record(&r, &validators[0]));
        assert_eq!(validators.len(), 1);
    }

    #[test]
    fn test_validator_inductor_status() {
        let validators = vec![];

        let r = get_registration();

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let _ = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let _ = inductor.induct(&r, ValidatorStatus::Active);
        let validators = inductor.to_vec();

        assert!(validators[0].status == ValidatorStatus::PendingActivation as u8);
        assert!(validators[1].status == ValidatorStatus::Active as u8);
        assert_eq!(validators.len(), 2);
    }

    #[test]
    fn test_validator_inductor_valid_all_active_validators() {
        let mut validators = vec![];
        for _ in 0..5 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Active as u8;
            validators.push(v);
        }

        let r = get_registration();

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(result.unwrap(), 5);
        assert!(registration_equals_record(&r, &validators[5]));
        assert_eq!(validators.len(), 6);
    }

    #[test]
    fn test_validator_inductor_valid_all_second_validator_withdrawn() {
        let mut validators = vec![];
        let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
        v.status = ValidatorStatus::Active as u8;
        validators.push(v);
        for _ in 0..4 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Withdrawn as u8;
            validators.push(v);
        }

        let r = get_registration();

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(result.unwrap(), 1);
        assert!(registration_equals_record(&r, &validators[1]));
        assert_eq!(validators.len(), 5);
    }

    #[test]
    fn test_validator_inductor_valid_all_withdrawn_validators() {
        let mut validators = vec![];
        for _ in 0..5 {
            let (mut v, _) = ValidatorRecord::zero_with_thread_rand_keypair();
            v.status = ValidatorStatus::Withdrawn as u8;
            validators.push(v);
        }

        /*
         * Ensure the first validator gets the 0'th slot
         */
        let r = get_registration();
        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();
        assert_eq!(result.unwrap(), 0);
        assert!(registration_equals_record(&r, &validators[0]));

        /*
         * Ensure the second validator gets the 1'st slot
         */
        let r_two = get_registration();
        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r_two, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();
        assert_eq!(result.unwrap(), 1);
        assert!(registration_equals_record(&r_two, &validators[1]));
        assert_eq!(validators.len(), 5);
    }

    #[test]
    fn test_validator_inductor_shard_too_high() {
        let validators = vec![];

        let mut r = get_registration();
        r.withdrawal_shard = 1025;

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(result, Err(ValidatorInductionError::InvalidShard));
        assert_eq!(validators.len(), 0);
    }

    #[test]
    fn test_validator_inductor_shard_proof_of_possession_failure() {
        let validators = vec![];

        let mut r = get_registration();
        let kp = Keypair::random();
        r.proof_of_possession = get_proof_of_possession(&kp);

        let mut inductor = ValidatorInductor::new(0, 1024, validators);
        let result = inductor.induct(&r, ValidatorStatus::PendingActivation);
        let validators = inductor.to_vec();

        assert_eq!(
            result,
            Err(ValidatorInductionError::InvaidProofOfPossession)
        );
        assert_eq!(validators.len(), 0);
    }
}
<|endoftext|>