<commit_before>use update_client::*;
use errors::*;

use crypto::digest::Digest;
use crypto::sha2::Sha256;
use base64::*;
use std;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};

pub trait Database: Send + Sync + Clone {
    fn query(&self, &str, ThreatDescriptor) -> Result<Option<()>>;
    fn update(&self, &FetchResponse) -> Result<()>;
    fn validate(&self, &FetchResponse) -> Result<()>;
}

#[derive(Clone)]
pub struct HashDB {
    inner_db: Arc<Mutex<HashMap<ThreatDescriptor, Vec<String>>>>,
}

impl HashDB {
    pub fn new() -> HashDB {
        HashDB { inner_db: Arc::new(Mutex::new(HashMap::new())) }
    }

    fn remove(&self,
              removal_map: &HashMap<ThreatDescriptor, (ResponseType, HashSet<usize>)>)
              -> Result<()> {

        for (descriptor, &(_response_type, ref removals)) in removal_map {
            if removals.is_empty() {
                continue;
            }

            info!("Removing {:#?} from {:#?}", removals.len(), descriptor);

            let mut cur_map = self.inner_db.lock().expect("Failed to attain lock for inner_db");

            let mut cur_hashes = cur_map.entry(*descriptor)
                                        .or_insert(vec![]);

            let new_hashes: Vec<String> = cur_hashes.iter()
                                                    .enumerate()
                                                    .filter_map(|(ix, s)| {
                                                        if removals.contains(&ix) {
                                                            Some(s.clone())
                                                        } else {
                                                            None
                                                        }
                                                    })
                                                    .collect();

            *cur_hashes = new_hashes;
        }

        Ok(())
    }

    fn add(&self,
           addition_map: &HashMap<ThreatDescriptor, (ResponseType, Vec<String>)>)
           -> Result<()> {

        for (descriptor, &(response_type, ref additions)) in addition_map {
            if additions.is_empty() {
                continue;
            }
            info!("Adding {:#?} to {:#?}", additions.len(), descriptor);
            let mut cur_map = self.inner_db.lock().expect("Failed to attain lock for inner_db");

            if let ResponseType::PartialUpdate = response_type {

                let mut cur_hashes = cur_map.entry(*descriptor)
                                            .or_insert(vec![]);

                cur_hashes.extend_from_slice(&additions[..]);

                cur_hashes.sort();
            } else if let ResponseType::FullUpdate = response_type {
                cur_map.insert(*descriptor, additions.clone());
            }
        }

        Ok(())
    }
}

impl Database for HashDB {
    fn update(&self, res: &FetchResponse) -> Result<()> {
        let rem_indices = try!(removals(res));
        let add_prefixes = try!(additions(res));

        try!(self.remove(&rem_indices));
        try!(self.add(&add_prefixes));

        Ok(())
    }

    fn validate(&self, res: &FetchResponse) -> Result<()> {
        for response in &res.list_update_responses {
            let descriptor = ThreatDescriptor {
                threat_type: response.threat_type,
                platform_type: response.platform_type,
                threat_entry_type: response.threat_entry_type,
            };
            let checksum = &response.checksum;

            let mut hash = {
                let mut cur_map = self.inner_db.lock().expect("Failed to attain lock for inner_db");

                let cur_hashes = match cur_map.get_mut(&descriptor) {
                    Some(h) => h,
                    None => {
                        info!("Could not get hashes for descriptor: {:#?}", descriptor);
                        continue;
                    }
                };

                cur_hashes.sort();
                cur_hashes.dedup();

                cur_hashes.iter()
                          .fold(Sha256::new(), |mut s, r| {
                              s.input(r.as_bytes());
                              s
                          })
            };

            let result = {
                hash.input(&vec![]);
                let mut result = vec![0; 32];
                hash.result(&mut result);
                result
            };
            let hash = result;

            if hash != checksum.sha256.as_bytes() {
                error!("Checksum failed {:?} != {:?}",
                       hash,
                       checksum.sha256.as_bytes());
            } else {
                info!("Database validated");
            }

        }

        Ok(())
    }

    fn query(&self, _url: &str, _descriptor: ThreatDescriptor) -> Result<Option<()>> {
        unimplemented!()
    }
}

fn additions(fetch_response: &FetchResponse)
             -> Result<HashMap<ThreatDescriptor, (ResponseType, Vec<String>)>> {

    let mut threat_map = HashMap::new();

    for response in &fetch_response.list_update_responses {
        let mut hash_prefixes = Vec::new();
        // Each set of additions will contain a string. This string is the concatenated list of
        // hash prefixes. In a single 'addition' threat entry the prefix size is the same for all
        // hashes, but the prefix size may vary across threat entries.
        for threat_entry in &response.additions {
            if let CompressionType::Raw = threat_entry.compression_type {

                let raw_hashes = &threat_entry.raw_hashes;

                if raw_hashes.raw_hashes.len() % raw_hashes.prefix_size as usize != 0 {
                    error!("Raw hashes not divisible by prefix_size. Skipping update.");
                    continue;
                }

                let hash_count = raw_hashes.raw_hashes.len() / raw_hashes.prefix_size as usize;
                hash_prefixes.reserve(hash_count);

                let init = raw_hashes.raw_hashes.as_bytes();
                let mut ix = 0;
                loop {
                    if ix == init.len() {
                        break;
                    }

                    let char_slice =
                        try!(String::from_utf8(init[ix..ix + raw_hashes.prefix_size as usize]
                                                   .to_vec())
                                 .chain_err(|| "could not convert hash to utf8 string"));
                    ix = ix + raw_hashes.prefix_size as usize;
                    hash_prefixes.push(char_slice);
                }
            }
        }

        threat_map.insert(ThreatDescriptor {
                              threat_type: response.threat_type,
                              platform_type: response.platform_type,
                              threat_entry_type: response.threat_entry_type,
                          },
                          (response.response_type, hash_prefixes));
    }
    Ok(threat_map)
}

fn removals(fetch_response: &FetchResponse)
            -> Result<HashMap<ThreatDescriptor, (ResponseType, HashSet<usize>)>> {

    let mut threat_map = HashMap::new();

    for response in &fetch_response.list_update_responses {
        let mut raw_indices = HashSet::new();

        for threat_entry in &response.removals {
            if let CompressionType::Raw = threat_entry.compression_type {

                for raw in threat_entry.raw_indices.indices.clone() {
                    raw_indices.insert(raw);
                }
            }
        }

        threat_map.insert(ThreatDescriptor {
                              threat_type: response.threat_type,
                              platform_type: response.platform_type,
                              threat_entry_type: response.threat_entry_type,
                          },
                          (response.response_type, raw_indices));
    }
    Ok(threat_map)
}

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

    #[test]
    fn test_update() {
        let db = HashDB::new();
        let res = fetch_req();

        db.update(&res);
    }

    #[test]
    fn test_additions() {
        let res = fetch_req();
        let addition_map = additions(&res).unwrap();

        assert_eq!(addition_map.get(&ThreatDescriptor {
                                   threat_type: ThreatType::Malware,
                                   platform_type: PlatformType::Linux,
                                   threat_entry_type: ThreatEntryType::Url,
                               })
                               .unwrap()
                               .1,
                   vec!["1234", "5678"]);
    }

    fn fetch_req() -> FetchResponse {

        let update_res = ListUpdateResponse {
            threat_type: ThreatType::Malware,
            threat_entry_type: ThreatEntryType::Url,
            platform_type: PlatformType::Linux,
            response_type: ResponseType::FullUpdate,
            additions: vec![ThreatEntrySet {
                                compression_type: CompressionType::Raw,
                                raw_hashes: RawHashes {
                                    prefix_size: 4,
                                    raw_hashes: "12345678".to_owned(),
                                },
                                raw_indices: RawIndices { indices: vec![] },
                                rice_hashes: RiceDeltaEncoding {
                                    first_value: "".to_owned(),
                                    rice_parameter: 0,
                                    num_entries: 0,
                                    encoded_data: "".to_owned(),
                                },
                                rice_indices: RiceDeltaEncoding {
                                    first_value: "".to_owned(),
                                    rice_parameter: 0,
                                    num_entries: 0,
                                    encoded_data: "".to_owned(),
                                },
                            }],
            removals: vec![],
            new_client_state: "new_state".to_owned(),
            checksum: Checksum { sha256: "".to_owned() },
        };

        FetchResponse {
            list_update_responses: vec![update_res],
            minimum_wait_duration: "300.4s".to_owned(),
        }
    }
    //

}
<commit_msg>Sort before removals.<commit_after>use update_client::*;
use errors::*;

use crypto::digest::Digest;
use crypto::sha2::Sha256;
use base64::*;
use std;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};

pub trait Database: Send + Sync + Clone {
    fn query(&self, &str, ThreatDescriptor) -> Result<Option<()>>;
    fn update(&self, &FetchResponse) -> Result<()>;
    fn validate(&self, &FetchResponse) -> Result<()>;
}

#[derive(Clone)]
pub struct HashDB {
    inner_db: Arc<Mutex<HashMap<ThreatDescriptor, Vec<String>>>>,
}

impl HashDB {
    pub fn new() -> HashDB {
        HashDB { inner_db: Arc::new(Mutex::new(HashMap::new())) }
    }

    fn remove(&self,
              removal_map: &HashMap<ThreatDescriptor, (ResponseType, HashSet<usize>)>)
              -> Result<()> {

        for (descriptor, &(_response_type, ref removals)) in removal_map {
            if removals.is_empty() {
                continue;
            }

            info!("Removing {:#?} from {:#?}", removals.len(), descriptor);

            let mut cur_map = self.inner_db.lock().expect("Failed to attain lock for inner_db");

            let mut cur_hashes = cur_map.entry(*descriptor)
                                        .or_insert(vec![]);

            cur_hashes.sort();
            let new_hashes: Vec<String> = cur_hashes.iter()
                                                    .enumerate()
                                                    .filter_map(|(ix, s)| {
                                                        if removals.contains(&ix) {
                                                            Some(s.clone())
                                                        } else {
                                                            None
                                                        }
                                                    })
                                                    .collect();

            *cur_hashes = new_hashes;
        }

        Ok(())
    }

    fn add(&self,
           addition_map: &HashMap<ThreatDescriptor, (ResponseType, Vec<String>)>)
           -> Result<()> {

        for (descriptor, &(response_type, ref additions)) in addition_map {
            if additions.is_empty() {
                continue;
            }
            info!("Adding {:#?} to {:#?}", additions.len(), descriptor);
            let mut cur_map = self.inner_db.lock().expect("Failed to attain lock for inner_db");

            if let ResponseType::PartialUpdate = response_type {

                let mut cur_hashes = cur_map.entry(*descriptor)
                                            .or_insert(vec![]);

                cur_hashes.extend_from_slice(&additions[..]);

                cur_hashes.sort();
            } else if let ResponseType::FullUpdate = response_type {
                cur_map.insert(*descriptor, additions.clone());
            }
        }

        Ok(())
    }
}

impl Database for HashDB {
    fn update(&self, res: &FetchResponse) -> Result<()> {
        let rem_indices = try!(removals(res));
        let add_prefixes = try!(additions(res));

        try!(self.remove(&rem_indices));
        try!(self.add(&add_prefixes));

        Ok(())
    }

    fn validate(&self, res: &FetchResponse) -> Result<()> {
        for response in &res.list_update_responses {
            let descriptor = ThreatDescriptor {
                threat_type: response.threat_type,
                platform_type: response.platform_type,
                threat_entry_type: response.threat_entry_type,
            };
            let checksum = &response.checksum;

            let mut hash = {
                let mut cur_map = self.inner_db.lock().expect("Failed to attain lock for inner_db");

                let cur_hashes = match cur_map.get_mut(&descriptor) {
                    Some(h) => h,
                    None => {
                        info!("Could not get hashes for descriptor: {:#?}", descriptor);
                        continue;
                    }
                };

                cur_hashes.sort();
                cur_hashes.dedup();

                cur_hashes.iter()
                          .fold(Sha256::new(), |mut s, r| {
                              s.input(r.as_bytes());
                              s
                          })
            };

            let result = {
                hash.input(&vec![]);
                let mut result = vec![0; 32];
                hash.result(&mut result);
                result
            };
            let hash = result;

            if hash != checksum.sha256.as_bytes() {
                error!("Checksum failed {:?} != {:?}",
                       hash,
                       checksum.sha256.as_bytes());
            } else {
                info!("Database validated");
            }

        }

        Ok(())
    }

    fn query(&self, _url: &str, _descriptor: ThreatDescriptor) -> Result<Option<()>> {
        unimplemented!()
    }
}

fn additions(fetch_response: &FetchResponse)
             -> Result<HashMap<ThreatDescriptor, (ResponseType, Vec<String>)>> {

    let mut threat_map = HashMap::new();

    for response in &fetch_response.list_update_responses {
        let mut hash_prefixes = Vec::new();
        // Each set of additions will contain a string. This string is the concatenated list of
        // hash prefixes. In a single 'addition' threat entry the prefix size is the same for all
        // hashes, but the prefix size may vary across threat entries.
        for threat_entry in &response.additions {
            if let CompressionType::Raw = threat_entry.compression_type {

                let raw_hashes = &threat_entry.raw_hashes;

                if raw_hashes.raw_hashes.len() % raw_hashes.prefix_size as usize != 0 {
                    error!("Raw hashes not divisible by prefix_size. Skipping update.");
                    continue;
                }

                let hash_count = raw_hashes.raw_hashes.len() / raw_hashes.prefix_size as usize;
                hash_prefixes.reserve(hash_count);

                let init = raw_hashes.raw_hashes.as_bytes();
                let mut ix = 0;
                loop {
                    if ix == init.len() {
                        break;
                    }

                    let char_slice =
                        try!(String::from_utf8(init[ix..ix + raw_hashes.prefix_size as usize]
                                                   .to_vec())
                                 .chain_err(|| "could not convert hash to utf8 string"));
                    ix = ix + raw_hashes.prefix_size as usize;
                    hash_prefixes.push(char_slice);
                }
            }
        }

        threat_map.insert(ThreatDescriptor {
                              threat_type: response.threat_type,
                              platform_type: response.platform_type,
                              threat_entry_type: response.threat_entry_type,
                          },
                          (response.response_type, hash_prefixes));
    }
    Ok(threat_map)
}

fn removals(fetch_response: &FetchResponse)
            -> Result<HashMap<ThreatDescriptor, (ResponseType, HashSet<usize>)>> {

    let mut threat_map = HashMap::new();

    for response in &fetch_response.list_update_responses {
        let mut raw_indices = HashSet::new();

        for threat_entry in &response.removals {
            if let CompressionType::Raw = threat_entry.compression_type {

                for raw in threat_entry.raw_indices.indices.clone() {
                    raw_indices.insert(raw);
                }
            }
        }

        threat_map.insert(ThreatDescriptor {
                              threat_type: response.threat_type,
                              platform_type: response.platform_type,
                              threat_entry_type: response.threat_entry_type,
                          },
                          (response.response_type, raw_indices));
    }
    Ok(threat_map)
}

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

    #[test]
    fn test_update() {
        let db = HashDB::new();
        let res = fetch_req();

        db.update(&res);
    }

    // Given an empty local database the database should recieve
    // a full update of hashes as well as a sha256 checksum
    // Assert that the checksum is as expected.

    #[test]
    fn test_additions() {
        let res = fetch_req();
        let addition_map = additions(&res).unwrap();

        assert_eq!(addition_map.get(&ThreatDescriptor {
                                   threat_type: ThreatType::Malware,
                                   platform_type: PlatformType::Linux,
                                   threat_entry_type: ThreatEntryType::Url,
                               })
                               .unwrap()
                               .1,
                   vec!["1234", "5678"]);
    }

    fn fetch_req() -> FetchResponse {

        let update_res = ListUpdateResponse {
            threat_type: ThreatType::Malware,
            threat_entry_type: ThreatEntryType::Url,
            platform_type: PlatformType::Linux,
            response_type: ResponseType::FullUpdate,
            additions: vec![ThreatEntrySet {
                                compression_type: CompressionType::Raw,
                                raw_hashes: RawHashes {
                                    prefix_size: 4,
                                    raw_hashes: "12345678".to_owned(),
                                },
                                raw_indices: RawIndices { indices: vec![] },
                                rice_hashes: RiceDeltaEncoding {
                                    first_value: "".to_owned(),
                                    rice_parameter: 0,
                                    num_entries: 0,
                                    encoded_data: "".to_owned(),
                                },
                                rice_indices: RiceDeltaEncoding {
                                    first_value: "".to_owned(),
                                    rice_parameter: 0,
                                    num_entries: 0,
                                    encoded_data: "".to_owned(),
                                },
                            }],
            removals: vec![],
            new_client_state: "new_state".to_owned(),
            checksum: Checksum { sha256: "".to_owned() },
        };

        FetchResponse {
            list_update_responses: vec![update_res],
            minimum_wait_duration: "300.4s".to_owned(),
        }
    }
    //

}
<|endoftext|>


<commit_before>//! This module provides an oct-tree of `Cube`s, which are the basic building
//! blocks of **Bismuth**.
//!
//! The tree is internally comprised of `Node`s. `BranchNode`s contain links to
//! sub-trees and `LeafNode`s contain geometric data associated with their
//! respective spatial partition. However, `Node`s do not contain any spatial
//! data (origin, partition, etc.); instead, this is calculated during
//! traversals by `Cube`, which acts as a recursive view over the `Node`s in the
//! tree and provides the primary interface for interacting with trees.
//!
//! Leaf `Cube`s provide the geometric and spatial data that together represent
//! the shape and layout of the game world.
//!
//! `Cube`s reference `Node`s in the tree, and abstract over the mutability of
//! those references. This module exposes the `CubeRef` and `CubeMut` type
//! definitions for immutable and mutable `Cube`s, respectively.
//!
//! Because `Cube`s reference `Node`s and `Node`s may reference other `Node`s
//! in a tree, orphan types are provided that only expose payloads and no links.
//! This is useful when collecting `Cube`s, because otherwise the references
//! into the rest of the tree would lead to aliasing. `OrphanCubeRef` and
//! `OrphanCubeMut` are the orphan analogues of `CubeRef` and `CubeMut`,
//! respectively. Orphans of course do not support indexing or traversal.
//!
//! `Root` can be used to create a new tree, and owns the root `Node`. `Root`s
//! expose `Cube`s to manipulate the tree.
//!
//! In the abstract, "cube" refers to the amalgamation of all the types used to
//! represent elements in a tree, which together form the complete notion of a
//! cube.
//!
//! # Examples
//!
//! Subdividing and iterating over the cubes in a tree:
//!
//! ```
//! use bismuth::cube::{LogWidth, Root, Spatial};
//!
//! let mut root = Root::new(LogWidth::max_value());
//! let _ = root.to_cube_mut().subdivide();
//! for cube in root.to_cube().iter() {
//!     println!("origin: {}; width: {}",
//!              cube.partition().origin(),
//!              cube.partition().width().to_inner());
//! }
//! ```

mod edit;
mod geometry;
mod mesh;
mod space;
#[macro_use]
mod traverse;
mod tree;

use self::tree::{Cube, OrphanCube};

pub use self::edit::Cursor;
pub use self::geometry::{Edge, Geometry, Offset};
pub use self::space::{AABB, Axis, Direction, Intersects, LogWidth, Orientation, Partition,
                      RayCast, RayIntersection, Spatial};
pub use self::tree::{BranchNode, BranchPayload, LeafNode, LeafPayload, Node, OrphanNode, Root};

pub type CubeRef<'a, 'b> = Cube<'a, &'b Node>;
pub type CubeMut<'a, 'b> = Cube<'a, &'b mut Node>;
pub type OrphanCubeRef<'a, 'b> = OrphanCube<'a, &'b LeafPayload, &'b BranchPayload>;
pub type OrphanCubeMut<'a, 'b> = OrphanCube<'a, &'b mut LeafPayload, &'b mut BranchPayload>;

#[cfg(test)]
mod tests {
    use super::*;
}
<commit_msg>Fix documentation code.<commit_after>//! This module provides an oct-tree of `Cube`s, which are the basic building
//! blocks of **Bismuth**.
//!
//! The tree is internally comprised of `Node`s. `BranchNode`s contain links to
//! sub-trees and `LeafNode`s contain geometric data associated with their
//! respective spatial partition. However, `Node`s do not contain any spatial
//! data (origin, partition, etc.); instead, this is calculated during
//! traversals by `Cube`, which acts as a recursive view over the `Node`s in the
//! tree and provides the primary interface for interacting with trees.
//!
//! Leaf `Cube`s provide the geometric and spatial data that together represent
//! the shape and layout of the game world.
//!
//! `Cube`s reference `Node`s in the tree, and abstract over the mutability of
//! those references. This module exposes the `CubeRef` and `CubeMut` type
//! definitions for immutable and mutable `Cube`s, respectively.
//!
//! Because `Cube`s reference `Node`s and `Node`s may reference other `Node`s
//! in a tree, orphan types are provided that only expose payloads and no links.
//! This is useful when collecting `Cube`s, because otherwise the references
//! into the rest of the tree would lead to aliasing. `OrphanCubeRef` and
//! `OrphanCubeMut` are the orphan analogues of `CubeRef` and `CubeMut`,
//! respectively. Orphans of course do not support indexing or traversal.
//!
//! `Root` can be used to create a new tree, and owns the root `Node`. `Root`s
//! expose `Cube`s to manipulate the tree.
//!
//! In the abstract, "cube" refers to the amalgamation of all the types used to
//! represent elements in a tree, which together form the complete notion of a
//! cube.
//!
//! # Examples
//!
//! Subdividing and iterating over the cubes in a tree:
//!
//! ```
//! use bismuth::cube::{LogWidth, Root, Spatial};
//!
//! let mut root = Root::new(LogWidth::max_value());
//! let _ = root.as_cube_mut().subdivide();
//! for cube in root.as_cube().iter() {
//!     println!("origin: {}; width: {}",
//!              cube.partition().origin(),
//!              cube.partition().width().to_inner());
//! }
//! ```

mod edit;
mod geometry;
mod mesh;
mod space;
#[macro_use]
mod traverse;
mod tree;

use self::tree::{Cube, OrphanCube};

pub use self::edit::Cursor;
pub use self::geometry::{Edge, Geometry, Offset};
pub use self::space::{AABB, Axis, Direction, Intersects, LogWidth, Orientation, Partition,
                      RayCast, RayIntersection, Spatial};
pub use self::tree::{BranchNode, BranchPayload, LeafNode, LeafPayload, Node, OrphanNode, Root};

pub type CubeRef<'a, 'b> = Cube<'a, &'b Node>;
pub type CubeMut<'a, 'b> = Cube<'a, &'b mut Node>;
pub type OrphanCubeRef<'a, 'b> = OrphanCube<'a, &'b LeafPayload, &'b BranchPayload>;
pub type OrphanCubeMut<'a, 'b> = OrphanCube<'a, &'b mut LeafPayload, &'b mut BranchPayload>;

#[cfg(test)]
mod tests {
    use super::*;
}
<|endoftext|>


<commit_before>use syntax::util::small_vector::SmallVector;

use syntax::{abi, ast, ast_util, codemap, parse};
use syntax::ptr::P;
use syntax::ext::base;
use syntax::ext::build::AstBuilder;

static BEFORE_EACH: &'static str = "before_each";
static AFTER_EACH:  &'static str = "after_each";
static BEFORE:      &'static str = "before";
static AFTER:       &'static str = "after";
static IT:          &'static str = "it";

/// Defines the state of a `describe!` macro as it is parsing.
struct DescribeState {
    name: Option<String>,
    before: Option<P<ast::Block>>,
    after: Option<P<ast::Block>>,
    before_each: Option<P<ast::Block>>,
    after_each: Option<P<ast::Block>>,
    tests: Vec<Test>
}

/// A test as a description and associated block.
struct Test {
    description: String,
    block: P<ast::Block>
}

fn parse_describe(mut state: DescribeState, mut parser: parse::parser::Parser) -> DescribeState {
    // First parse the name of this describe block:
    let (name, _) = parser.parse_str();
    state.name = Some(name.get().to_string());

    // Move past the opening {
    if parse::token::LBRACE != parser.bump_and_get() {
        parser.fatal("Expected { after the description of a describe! block.");
    }

    // Now parse all tests and subsections:
    while parser.token != parse::token::RBRACE {
        // Get the name of this block, must be either:
        //     - before_each
        //     - after_each
        //     - before
        //     - after
        //     - it
        //
        // Any other top-level idents are not allowed.
        let block_name = parser.parse_ident();

        match block_name.as_str() {
            BEFORE_EACH => {
                if state.before_each.is_some() { parser.fatal("Only one `before_each` block is allowed per `describe!` block.") }
                state.before_each = Some(parser.parse_block());
            },

            AFTER_EACH => {
                if state.after_each.is_some() { parser.fatal("Only one `after_each` block is allowed per `describe!` block.") }
                state.after_each = Some(parser.parse_block());
            },

            BEFORE => {
                if state.before.is_some() { parser.fatal("Only one `before` block is allowed per `describe!` block.") }
                state.before = Some(parser.parse_block());
            },

            AFTER => {
                if state.after.is_some() { parser.fatal("Only one `after` block is allowed per `describe!` block.") }
                state.after = Some(parser.parse_block());
            },

            IT => {
                // Description of this `it` block.
                let (description, _) = parser.parse_str();

                state.tests.push(Test {
                    // Get as a String
                    description: description.get().to_string(),

                    // The associated block
                    block: parser.parse_block()
                });
            },

            banned => {
                // Illegal block name.
                let span = parser.span;
                parser.span_fatal(span, format!("Expected one of: `{}`, but found: `{}`",
                    format!("{}, {}, {}, {}, {}", BEFORE_EACH, AFTER_EACH, BEFORE, AFTER, IT).as_slice(),
                    banned).as_slice());
            }
        }
    }

    state
}

<commit_msg>(feat) Test creator (no support for before/after yet).<commit_after>use syntax::util::small_vector::SmallVector;

use syntax::{abi, ast, ast_util, codemap, parse};
use syntax::ptr::P;
use syntax::ext::base;
use syntax::ext::build::AstBuilder;

static BEFORE_EACH: &'static str = "before_each";
static AFTER_EACH:  &'static str = "after_each";
static BEFORE:      &'static str = "before";
static AFTER:       &'static str = "after";
static IT:          &'static str = "it";

/// Defines the state of a `describe!` macro as it is parsing.
struct DescribeState {
    name: Option<String>,
    before: Option<P<ast::Block>>,
    after: Option<P<ast::Block>>,
    before_each: Option<P<ast::Block>>,
    after_each: Option<P<ast::Block>>,
    tests: Vec<Test>
}

/// A test as a description and associated block.
struct Test {
    description: String,
    block: P<ast::Block>
}

fn create_tests(state: DescribeState, cx: &mut base::ExtCtxt) -> Vec<P<ast::Item>> {
    let name = state.name.unwrap().replace(" ", "_");

    // FIXME(reem): Implement before and after.
    let (_before, _after) = (state.before, state.after);

    let (before_each, after_each) = (state.before_each, state.after_each);
    let tests = state.tests;

    let test_attribute = cx.attribute(codemap::DUMMY_SP,
                                      cx.meta_word(codemap::DUMMY_SP, parse::token::InternedString::new("test")));

    tests.into_iter().map(|Test { description, block }| {
        let test_body = match (&before_each, &after_each) {
            (&None, &None) => block,

            (&Some(ref before), &None) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items,
                    stmts: before.stmts + block.stmts,
                    ..block.deref().clone()
                })
            },

            (&None, &Some(ref after)) => {
                P(ast::Block {
                    view_items: block.view_items + after.view_items,
                    stmts: block.stmts + after.stmts,
                    ..block.deref().clone()
                })
            },

            (&Some(ref before), &Some(ref after)) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items + after.view_items,
                    stmts: before.stmts + block.stmts + after.stmts,
                    ..block.deref().clone()
                })
            }
        };

        let mut test_name = name.clone();
        test_name.push_str("_");
        test_name.push_str(description.replace(" ", "_").as_slice());

        P(ast::Item {
            ident: cx.ident_of(test_name.as_slice()),
            attrs: vec![test_attribute.clone()],
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemFn(
                cx.fn_decl(vec![], cx.ty_nil()),
                ast::NormalFn,
                abi::Rust,
                ast_util::empty_generics(),
                test_body
            ),
            vis: ast::Inherited,
            span: codemap::DUMMY_SP
        })
    }).collect()
}

fn parse_describe(mut state: DescribeState, mut parser: parse::parser::Parser) -> DescribeState {
    // First parse the name of this describe block:
    let (name, _) = parser.parse_str();
    state.name = Some(name.get().to_string());

    // Move past the opening {
    if parse::token::LBRACE != parser.bump_and_get() {
        parser.fatal("Expected { after the description of a describe! block.");
    }

    // Now parse all tests and subsections:
    while parser.token != parse::token::RBRACE {
        // Get the name of this block, must be either:
        //     - before_each
        //     - after_each
        //     - before
        //     - after
        //     - it
        //
        // Any other top-level idents are not allowed.
        let block_name = parser.parse_ident();

        match block_name.as_str() {
            BEFORE_EACH => {
                if state.before_each.is_some() { parser.fatal("Only one `before_each` block is allowed per `describe!` block.") }
                state.before_each = Some(parser.parse_block());
            },

            AFTER_EACH => {
                if state.after_each.is_some() { parser.fatal("Only one `after_each` block is allowed per `describe!` block.") }
                state.after_each = Some(parser.parse_block());
            },

            BEFORE => {
                if state.before.is_some() { parser.fatal("Only one `before` block is allowed per `describe!` block.") }
                state.before = Some(parser.parse_block());
            },

            AFTER => {
                if state.after.is_some() { parser.fatal("Only one `after` block is allowed per `describe!` block.") }
                state.after = Some(parser.parse_block());
            },

            IT => {
                // Description of this `it` block.
                let (description, _) = parser.parse_str();

                state.tests.push(Test {
                    // Get as a String
                    description: description.get().to_string(),

                    // The associated block
                    block: parser.parse_block()
                });
            },

            banned => {
                // Illegal block name.
                let span = parser.span;
                parser.span_fatal(span, format!("Expected one of: `{}`, but found: `{}`",
                    format!("{}, {}, {}, {}, {}", BEFORE_EACH, AFTER_EACH, BEFORE, AFTER, IT).as_slice(),
                    banned).as_slice());
            }
        }
    }

    state
}

<|endoftext|>


<commit_before>use cargo_helper::CargoInfo;
use Bytes;
use client::GetResponse;
use hyper::client::Client;
use hyper::error::Error;
use hyper::header::{ByteRangeSpec, Headers, Range};
use pbr::{MultiBar, Pipe, ProgressBar, Units};
use response::CheckResponseStatus;
use std::cmp::min;
use std::io::Read;
use std::thread;
use std::time::{Instant, Duration};
use write::{OutputFileWriter, OutputChunkWriter};

/// Constant to represent the length of the buffer to download
/// the remote content
const DOWNLOAD_BUFFER_BYTES: usize = 1024 * 64;

/// Constant to represent the refresh interval (in milliseconds)
/// for the CLI
const PROGRESS_UPDATE_INTERVAL_MILLIS: u64 = 500;

/// Represents a range between two Bytes types
#[derive(Debug, PartialEq)]
struct RangeBytes(Bytes, Bytes);

macro_rules! initbar {
    ($mp:ident,$mpb:ident,$length:expr,$index:expr) => {
        let mut $mp = $mpb.create_bar($length);
        $mp.tick_format("▏▎▍▌▋▊▉██▉▊▋▌▍▎▏");
        $mp.format("|#--|");
        $mp.show_tick = true;
        $mp.show_speed = true;
        $mp.show_percent = true;
        $mp.show_counter = false;
        $mp.show_time_left = true;
        $mp.set_units(Units::Bytes);
        $mp.message(&format!("Chunk {} ", $index));
    }
}

/// Function to get the current chunk length, based on the chunk index.
fn get_chunk_length(chunk_index: u64,
                    content_length: Bytes,
                    global_chunk_length: Bytes)
                    -> Option<RangeBytes> {

    if content_length == 0 || global_chunk_length == 0 {
        return None;
    }

    let b_range: Bytes = chunk_index * global_chunk_length;

    if b_range >= (content_length - 1) {
        return None;
    }

    let e_range: Bytes = min(content_length - 1,
                             ((chunk_index + 1) * global_chunk_length) - 1);

    Some(RangeBytes(b_range, e_range))

}

/// Function to get the HTTP header to send to the file server, for a chunk (specified by its index)
fn get_header_from_index(chunk_index: u64,
                         content_length: Bytes,
                         global_chunk_length: Bytes)
                         -> Option<(Headers, RangeBytes)> {

    get_chunk_length(chunk_index, content_length, global_chunk_length).map(|range| {
        let mut header = Headers::new();
        header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(range.0, range.1)]));
        (header, RangeBytes(range.0, range.1 - range.0))
    })
}

/// Function to get from the server the content of a chunk.
/// This function returns a Result type - Bytes if the content of the header is accessible, an Error type otherwise.
fn download_a_chunk(http_client: &Client,
                    http_header: Headers,
                    mut chunk_writer: OutputChunkWriter,
                    url: &str,
                    mpb: &mut ProgressBar<Pipe>)
                    -> Result<Bytes, Error> {

    let mut body = http_client
        .get_http_response_using_headers(url, http_header)
        .unwrap();
    if !body.check_partialcontent_status() {
        return Err(Error::Status);
    }
    let mut bytes_buffer = [0; DOWNLOAD_BUFFER_BYTES];
    let mut sum_bytes = 0;

    let progress_update_interval = Duration::from_millis(PROGRESS_UPDATE_INTERVAL_MILLIS);
    let mut last_progress_bytes = 0;
    let mut last_progress_time = Instant::now() - progress_update_interval;

    while let Ok(n) = body.read(&mut bytes_buffer) {
        if n == 0 {
            return Ok(sum_bytes);
        }

        chunk_writer.write(sum_bytes, &bytes_buffer[0..n]);

        sum_bytes += n as u64;

        // Update the CLI
        if Instant::now().duration_since(last_progress_time) > progress_update_interval {
            last_progress_time = Instant::now();
            let progress_bytes_delta = sum_bytes - last_progress_bytes;
            last_progress_bytes = sum_bytes;
            mpb.add(progress_bytes_delta);
        }
    }
    mpb.add(sum_bytes - last_progress_bytes);
    return Ok(0u64);
}

/// Function to download each chunk of a remote content (given by its URL).
/// This function takes as parameters:
/// * the remote content length,
/// * a mutable reference to share between threads, which contains each chunk,
/// * the number of chunks that contains the remote content,
/// * the URL of the remote content server,
/// * a custom authorization to access and download the remote content.
pub fn download_chunks(cargo_info: CargoInfo,
                       mut out_file: OutputFileWriter,
                       nb_chunks: u64,
                       url: &str)
                       -> bool {
    let (content_length, auth_header_factory) = (cargo_info.content_length, cargo_info.auth_header);

    let global_chunk_length: u64 = (content_length / nb_chunks) + 1;
    let mut jobs = vec![];

    let mut mpb = MultiBar::new();
    mpb.println(&format!("Downloading {} chunks: ", nb_chunks));

    for chunk_index in 0..nb_chunks {

        let (mut http_header, RangeBytes(chunk_offset, chunk_length)) =
            get_header_from_index(chunk_index, content_length, global_chunk_length).unwrap();
        let hyper_client = Client::new();
        let url_clone = String::from(url);
        if let Some(auth_header_factory) = auth_header_factory.clone() {
            http_header.set(auth_header_factory.build_header());
        }

        // Initialize the progress bar for that chunk
        initbar!(mp, mpb, chunk_length, chunk_index);

        let chunk_writer = out_file.get_chunk_writer(chunk_offset);

        // In this work, we push a boolean value to know if the chunk is OK
        jobs.push(thread::spawn(move || match download_a_chunk(&hyper_client,
                                                               http_header,
                                                               chunk_writer,
                                                               &url_clone,
                                                               &mut mp) {
                                    Ok(bytes_written) => {
            mp.finish();
            if bytes_written == 0 {
                println!("The downloaded chunk {} is empty", chunk_index);
            }
            return true;
        }
                                    Err(error) => {
            mp.finish();
            println!("Cannot download the chunk {}, due to error {}",
                     chunk_index,
                     error);
            return false;
        }
                                }));
    }

    mpb.listen();

    // Contain the result state for chunks
    let mut child_results: Vec<bool> = Vec::with_capacity(nb_chunks as usize);

    for child in jobs {
        match child.join() {
            Ok(b) => child_results.push(b),
            Err(_) => child_results.push(false),
        }
    }

    // Check if all chunks are ok
    return child_results.iter().all(|x| *x);
}

#[cfg(test)]
mod test_chunk_length {

    use super::get_chunk_length;
    use super::RangeBytes;

    #[test]
    fn wrong_content_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 15, 0));
    }

    #[test]
    fn wrong_global_chunk_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 15));
    }

    #[test]
    fn wrong_length_parameters_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 0));
    }

    #[test]
    fn get_the_first_range_in_chunk() {
        assert_eq!(Some(RangeBytes(0, 249)), get_chunk_length(0, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_chunk() {
        assert_eq!(Some(RangeBytes(750, 999)), get_chunk_length(3, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_shorten_chunk() {
        assert_eq!(Some(RangeBytes(750, 997)), get_chunk_length(3, 998, 250));
    }

    #[test]
    fn wrong_index_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(4, 1000, 250));
    }

}

#[cfg(test)]
mod test_header {

    use super::{get_header_from_index, RangeBytes};
    use hyper::header::{ByteRangeSpec, Headers, Range};

    #[test]
    fn wrong_chunk_length_should_return_none() {
        assert_eq!(None, get_header_from_index(0, 0, 0));
    }

    #[test]
    fn good_chunk_length_should_return_a_good_header() {
        let mut test_header = Headers::new();
        test_header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(750, 997)]));
        assert_eq!(Some((test_header, RangeBytes(750, 247))),
                   get_header_from_index(3, 998, 250));
    }

}
<commit_msg>Fix typo<commit_after>use cargo_helper::CargoInfo;
use Bytes;
use client::GetResponse;
use hyper::client::Client;
use hyper::error::Error;
use hyper::header::{ByteRangeSpec, Headers, Range};
use pbr::{MultiBar, Pipe, ProgressBar, Units};
use response::CheckResponseStatus;
use std::cmp::min;
use std::io::Read;
use std::thread;
use std::time::{Instant, Duration};
use write::{OutputFileWriter, OutputChunkWriter};

/// Constant to represent the length of the buffer to download
/// the remote content
const DOWNLOAD_BUFFER_BYTES: usize = 1024 * 64;

/// Constant to represent the refresh interval (in milliseconds)
/// for the CLI
const PROGRESS_UPDATE_INTERVAL_MILLIS: u64 = 500;

/// Represents a range between two Bytes types
#[derive(Debug, PartialEq)]
struct RangeBytes(Bytes, Bytes);

macro_rules! initbar {
    ($mp:ident,$mpb:ident,$length:expr,$index:expr) => {
        let mut $mp = $mpb.create_bar($length);
        $mp.tick_format("▏▎▍▌▋▊▉██▉▊▋▌▍▎▏");
        $mp.format("|#--|");
        $mp.show_tick = true;
        $mp.show_speed = true;
        $mp.show_percent = true;
        $mp.show_counter = false;
        $mp.show_time_left = true;
        $mp.set_units(Units::Bytes);
        $mp.message(&format!("Chunk {} ", $index));
    }
}

/// Function to get the current chunk length, based on the chunk index.
fn get_chunk_length(chunk_index: u64,
                    content_length: Bytes,
                    global_chunk_length: Bytes)
                    -> Option<RangeBytes> {

    if content_length == 0 || global_chunk_length == 0 {
        return None;
    }

    let b_range: Bytes = chunk_index * global_chunk_length;

    if b_range >= (content_length - 1) {
        return None;
    }

    let e_range: Bytes = min(content_length - 1,
                             ((chunk_index + 1) * global_chunk_length) - 1);

    Some(RangeBytes(b_range, e_range))

}

/// Function to get the HTTP header to send to the file server, for a chunk (specified by its index)
fn get_header_from_index(chunk_index: u64,
                         content_length: Bytes,
                         global_chunk_length: Bytes)
                         -> Option<(Headers, RangeBytes)> {

    get_chunk_length(chunk_index, content_length, global_chunk_length).map(|range| {
        let mut header = Headers::new();
        header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(range.0, range.1)]));
        (header, RangeBytes(range.0, range.1 - range.0))
    })
}

/// Function to get from the server the content of a chunk.
/// This function returns a Result type - Bytes if the content of the header is accessible, an Error type otherwise.
fn download_a_chunk(http_client: &Client,
                    http_header: Headers,
                    mut chunk_writer: OutputChunkWriter,
                    url: &str,
                    mpb: &mut ProgressBar<Pipe>)
                    -> Result<Bytes, Error> {

    let mut body = http_client
        .get_http_response_using_headers(url, http_header)
        .unwrap();
    if !body.check_partialcontent_status() {
        return Err(Error::Status);
    }
    let mut bytes_buffer = [0; DOWNLOAD_BUFFER_BYTES];
    let mut sum_bytes = 0;

    let progress_update_interval = Duration::from_millis(PROGRESS_UPDATE_INTERVAL_MILLIS);
    let mut last_progress_bytes = 0;
    let mut last_progress_time = Instant::now() - progress_update_interval;

    while let Ok(n) = body.read(&mut bytes_buffer) {
        if n == 0 {
            return Ok(sum_bytes);
        }

        chunk_writer.write(sum_bytes, &bytes_buffer[0..n]);

        sum_bytes += n as u64;

        // Update the CLI
        if Instant::now().duration_since(last_progress_time) > progress_update_interval {
            last_progress_time = Instant::now();
            let progress_bytes_delta = sum_bytes - last_progress_bytes;
            last_progress_bytes = sum_bytes;
            mpb.add(progress_bytes_delta);
        }
    }
    mpb.add(sum_bytes - last_progress_bytes);
    return Ok(0u64);
}

/// Function to download each chunk of a remote content (given by its URL).
/// This function takes as parameters:
/// * the remote content length,
/// * a mutable reference to share between threads, which contains each chunk,
/// * the number of chunks that contains the remote content,
/// * the URL of the remote content server,
/// * a custom authorization to access and download the remote content.
pub fn download_chunks(cargo_info: CargoInfo,
                       mut out_file: OutputFileWriter,
                       nb_chunks: u64,
                       url: &str)
                       -> bool {
    let (content_length, auth_header_factory) = (cargo_info.content_length, cargo_info.auth_header);

    let global_chunk_length: u64 = (content_length / nb_chunks) + 1;
    let mut jobs = vec![];

    let mut mpb = MultiBar::new();
    mpb.println(&format!("Downloading {} chunks: ", nb_chunks));

    for chunk_index in 0..nb_chunks {

        let (mut http_header, RangeBytes(chunk_offset, chunk_length)) =
            get_header_from_index(chunk_index, content_length, global_chunk_length).unwrap();
        let hyper_client = Client::new();
        let url_clone = String::from(url);
        if let Some(auth_header_factory) = auth_header_factory.clone() {
            http_header.set(auth_header_factory.build_header());
        }

        // Initialize the progress bar for that chunk
        initbar!(mp, mpb, chunk_length, chunk_index);

        let chunk_writer = out_file.get_chunk_writer(chunk_offset);

        // In this work, we push a boolean value to know if the chunk is OK
        jobs.push(thread::spawn(move || match download_a_chunk(&hyper_client,
                                                               http_header,
                                                               chunk_writer,
                                                               &url_clone,
                                                               &mut mp) {
                                    Ok(bytes_written) => {
            mp.finish();
            if bytes_written == 0 {
                println!("The downloaded chunk {} is empty", chunk_index);
            }
            return true;
        }
                                    Err(error) => {
            mp.finish();
            println!("Cannot download the chunk {}, due to error {}",
                     chunk_index,
                     error);
            return false;
        }
                                }));
    }

    mpb.listen();

    // Contain the result state for chunks
    let mut child_results: Vec<bool> = Vec::with_capacity(nb_chunks as usize);

    for child in jobs {
        match child.join() {
            Ok(b) => child_results.push(b),
            Err(_) => child_results.push(false),
        }
    }

    // Check if all chunks are OK
    return child_results.iter().all(|x| *x);
}

#[cfg(test)]
mod test_chunk_length {

    use super::get_chunk_length;
    use super::RangeBytes;

    #[test]
    fn wrong_content_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 15, 0));
    }

    #[test]
    fn wrong_global_chunk_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 15));
    }

    #[test]
    fn wrong_length_parameters_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 0));
    }

    #[test]
    fn get_the_first_range_in_chunk() {
        assert_eq!(Some(RangeBytes(0, 249)), get_chunk_length(0, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_chunk() {
        assert_eq!(Some(RangeBytes(750, 999)), get_chunk_length(3, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_shorten_chunk() {
        assert_eq!(Some(RangeBytes(750, 997)), get_chunk_length(3, 998, 250));
    }

    #[test]
    fn wrong_index_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(4, 1000, 250));
    }

}

#[cfg(test)]
mod test_header {

    use super::{get_header_from_index, RangeBytes};
    use hyper::header::{ByteRangeSpec, Headers, Range};

    #[test]
    fn wrong_chunk_length_should_return_none() {
        assert_eq!(None, get_header_from_index(0, 0, 0));
    }

    #[test]
    fn good_chunk_length_should_return_a_good_header() {
        let mut test_header = Headers::new();
        test_header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(750, 997)]));
        assert_eq!(Some((test_header, RangeBytes(750, 247))),
                   get_header_from_index(3, 998, 250));
    }

}
<|endoftext|>


<commit_before>use bit_field::BitField;

use instruction::Op;
use register::Reg;
use register::Apsr;
use core::Core;
use memory::Fetch;

use operation::sign_extend;
use operation::add_with_carry;
use operation::condition_passed;

pub fn execute<T: Fetch>(core: &mut Core<T>, op: Option<Op>) {
    match op {
        None => panic!("undefined code"),
        Some(oper) => {
            match oper {
                Op::MOV { rd, rm } => {
                    core.r[rd.value() as usize] = core.r[rm.value()];
                    core.r[Reg::PC.value()] += 2;
                }
                Op::BL { imm32 } => {
                    let pc = core.r[Reg::PC.value()] + 4;
                    core.r[Reg::LR.value()] = pc | 0x01;
                    core.r[Reg::PC.value()] = ((pc as i32) + imm32) as u32;
                }
                Op::BX { rm } => {
                    core.r[Reg::PC.value()] = core.r[rm.value() as usize] & 0xfffffffe;
                }
                Op::BLX { rm } => {
                    let pc = core.r[Reg::PC.value()] + 4;
                    core.r[Reg::LR.value()] = (pc - 2) | 0x01;
                    core.r[Reg::PC.value()] = ((pc as i32) + (rm.value() as i32)) as u32;
                }
                Op::MOV_imm8 { rd, imm8 } => {
                    core.r[rd.value()] = imm8 as u32;
                    core.r[Reg::PC.value()] += 2;
                }
                Op::B_imm8 { cond, imm8 } => {
                    let imm32 = sign_extend((imm8 as u32) << 1, 8, 32);
                    if condition_passed(cond, &core.apsr) {
                        let pc = core.r[Reg::PC.value()] + 4;
                        core.r[Reg::PC.value()] = ((pc as i32) + imm32) as u32;
                    } else {
                        core.r[Reg::PC.value()] += 2;
                    }
                }
                Op::B_imm11 { imm11 } => {
                    let pc = core.r[Reg::PC.value()] + 4;
                    let imm32 = sign_extend((imm11 as u32) << 1, 11, 32);
                    core.r[Reg::PC.value()] = ((pc as i32) + imm32) as u32;
                }

                Op::CMP_imm8 { rn, imm8 } => {
                    let imm32 = imm8 as u32;
                    let (result, carry, overflow) =
                        add_with_carry(core.r[rn.value()], imm32 ^ 0xFFFFFFFF, true);
                    core.apsr.set_n(result.get_bit(31));
                    core.apsr.set_z(result == 0);
                    core.apsr.set_c(carry);
                    core.apsr.set_v(overflow);
                    core.r[Reg::PC.value()] += 2;
                }
                Op::CMP { rn, rm } => {
                    let (result, carry, overflow) =
                        add_with_carry(core.r[rn.value()], core.r[rm.value()] ^ 0xFFFFFFFF, true);
                    core.apsr.set_n(result.get_bit(31));
                    core.apsr.set_z(result == 0);
                    core.apsr.set_c(carry);
                    core.apsr.set_v(overflow);
                    core.r[Reg::PC.value()] += 2;
                }

                Op::PUSH { registers } => {
                    let address = core.msp - 4 * (registers.len() as u32);
                    core.r[Reg::PC.value()] += 2;
                }
                Op::LDR { rt, imm8 } => {
                    let imm32 = (imm8 as u32) << 2;
                    let address = (core.r[Reg::PC.value()] & 0xfffffffc) + imm32;
                    core.r[rt.value()] = core.memory.fetch32(address);
                    core.r[Reg::PC.value()] += 2;
                }
                Op::ADD { rdn, rm } => {
                    //TODO: reading R[15] should yield PC+4 here
                    let (result, carry, overflow) =
                        add_with_carry(core.r[rdn.value()], core.r[rm.value()], false);
                    core.r[rdn.value()] = result;
                    core.r[Reg::PC.value()] += 2;
                }
                Op::ADDS_imm { rn, rd, imm32 } => {
                    let (result, carry, overflow) =
                        add_with_carry(core.r[rn.value()], imm32 as u32, false);
                    core.apsr.set_n(result.get_bit(31));
                    core.apsr.set_z(result == 0);
                    core.apsr.set_c(carry);
                    core.apsr.set_v(overflow);

                    core.r[rd.value()] = result;
                    core.r[Reg::PC.value()] += 2;

                }
                Op::ADDS { rm, rn, rd } => {
                    let (result, carry, overflow) =
                        add_with_carry(core.r[rn.value()], core.r[rm.value()], false);
                    core.apsr.set_n(result.get_bit(31));
                    core.apsr.set_z(result == 0);
                    core.apsr.set_c(carry);
                    core.apsr.set_v(overflow);

                    core.r[rd.value()] = result;
                    core.r[Reg::PC.value()] += 2;

                }

                _ => panic!("unimplemented instruction"),
            }
        }
    }
}
<commit_msg>PC reg reading fix<commit_after>use bit_field::BitField;

use instruction::Op;
use register::Reg;
use register::Apsr;
use core::Core;
use memory::Fetch;

use operation::sign_extend;
use operation::add_with_carry;
use operation::condition_passed;

fn read_reg<T: Fetch>(core: &mut Core<T>, r: Reg) -> u32 {
    match r {
        Reg::PC => core.r[r.value()] + 4,
        _ => core.r[r.value()],
    }
}

pub fn execute<T: Fetch>(core: &mut Core<T>, op: Option<Op>) {
    match op {
        None => panic!("undefined code"),
        Some(oper) => {
            match oper {
                Op::MOV { rd, rm } => {
                    println!("MOV");
                    core.r[rd.value() as usize] = read_reg(core, rm);
                    core.r[Reg::PC.value()] += 2;
                }
                Op::BL { imm32 } => {
                    println!("BL");
                    let pc = read_reg(core, Reg::PC);
                    core.r[Reg::LR.value()] = pc | 0x01;
                    core.r[Reg::PC.value()] = ((pc as i32) + imm32) as u32;
                }
                Op::BX { rm } => {
                    println!("BX");
                    core.r[Reg::PC.value()] = read_reg(core, rm) & 0xfffffffe;
                }
                Op::BLX { rm } => {
                    println!("BLX");
                    let pc = read_reg(core, Reg::PC);
                    core.r[Reg::LR.value()] = (pc - 2) | 0x01;
                    core.r[Reg::PC.value()] = ((pc as i32) + (rm.value() as i32)) as u32;
                }
                Op::MOV_imm8 { rd, imm8 } => {
                    println!("MOV_imm8");
                    core.r[rd.value()] = imm8 as u32;
                    core.r[Reg::PC.value()] += 2;
                }
                Op::B_imm8 { cond, imm8 } => {
                    println!("B_imm8");
                    let imm32 = sign_extend((imm8 as u32) << 1, 8, 32);
                    if condition_passed(cond, &core.apsr) {
                        let pc = read_reg(core, Reg::PC);
                        core.r[Reg::PC.value()] = ((pc as i32) + imm32) as u32;
                    } else {
                        core.r[Reg::PC.value()] += 2;
                    }
                }
                Op::B_imm11 { imm11 } => {
                    println!("B_imm11");
                    let pc = read_reg(core, Reg::PC);
                    let imm32 = sign_extend((imm11 as u32) << 1, 11, 32);
                    core.r[Reg::PC.value()] = ((pc as i32) + imm32) as u32;
                }

                Op::CMP_imm8 { rn, imm8 } => {
                    println!("CMP_imm8");
                    let imm32 = imm8 as u32;
                    let (result, carry, overflow) =
                        add_with_carry(read_reg(core, rn), imm32 ^ 0xFFFFFFFF, true);
                    core.apsr.set_n(result.get_bit(31));
                    core.apsr.set_z(result == 0);
                    core.apsr.set_c(carry);
                    core.apsr.set_v(overflow);
                    core.r[Reg::PC.value()] += 2;
                }
                Op::CMP { rn, rm } => {
                    println!("CMP");
                    let (result, carry, overflow) =
                        add_with_carry(read_reg(core, rn), read_reg(core, rm) ^ 0xFFFFFFFF, true);
                    core.apsr.set_n(result.get_bit(31));
                    core.apsr.set_z(result == 0);
                    core.apsr.set_c(carry);
                    core.apsr.set_v(overflow);
                    core.r[Reg::PC.value()] += 2;
                }

                Op::PUSH { registers } => {
                    println!("PUSH");
                    let address = core.msp - 4 * (registers.len() as u32);
                    core.r[Reg::PC.value()] += 2;
                }
                Op::LDR { rt, imm8 } => {
                    println!("LDR");
                    let imm32 = (imm8 as u32) << 2;
                    let pc = read_reg(core, Reg::PC);
                    let address = (pc & 0xfffffffc) + imm32;
                    core.r[rt.value()] = core.memory.fetch32(address);
                    core.r[Reg::PC.value()] += 2;
                }
                Op::ADD { rdn, rm } => {

                    println!("ADD");
                    let (result, carry, overflow) =
                        add_with_carry(read_reg(core, rdn), read_reg(core, rm), false);
                    core.r[rdn.value()] = result;
                    core.r[Reg::PC.value()] += 2;
                }
                Op::ADDS_imm { rn, rd, imm32 } => {
                    println!("ADDS_imm R{:x}, R{:x}, #{:x}", rn.value(), rd.value(), imm32);
                    let r_n = read_reg(core, rn);
                    let (result, carry, overflow) =
                        add_with_carry(read_reg(core, rn), imm32 as u32, false);

                    core.apsr.set_n(result.get_bit(31));
                    core.apsr.set_z(result == 0);
                    core.apsr.set_c(carry);
                    core.apsr.set_v(overflow);

                    core.r[rd.value()] = result;
                    core.r[Reg::PC.value()] += 2;

                }
                Op::ADDS { rm, rn, rd } => {
                    println!("ADDS");
                    let (result, carry, overflow) =
                        add_with_carry(read_reg(core, rn), read_reg(core, rm), false);
                    core.apsr.set_n(result.get_bit(31));
                    core.apsr.set_z(result == 0);
                    core.apsr.set_c(carry);
                    core.apsr.set_v(overflow);

                    core.r[rd.value()] = result;
                    core.r[Reg::PC.value()] += 2;

                }

                _ => panic!("unimplemented instruction") ,
            }
        }
    }
}
<|endoftext|>


<commit_before>// mod escaping

const ESCAPE_CHARS: [(char, &'static str); 11] = [('\\', "\\\\"),
                                                  (' ', "\\s"),
                                                  ('/', "\\/"),
                                                  ('|', "\\p"),
                                                  (7 as char, "\\a"),
                                                  (8 as char, "\\b"),
                                                  (9 as char, "\\t"),
                                                  (10 as char, "\\n"),
                                                  (11 as char, "\\v"),
                                                  (12 as char, "\\f"),
                                                  (13 as char, "\\r")];

/// escapes all chars described in the server query manual
/// # Example
/// ```
/// use sqlib::escaping::escape;
///
/// let unescaped = "hello world/|\\".to_string();
/// let escaped = "hello\\sworld\\/\\p\\\\".to_string();
///
/// let s = escape(&unescaped);
///
/// assert_eq!(s, escaped);
/// ```
pub fn escape(s: &str) -> String {
    let mut new_string = s.to_string();
    for &(from, to) in ESCAPE_CHARS.iter() {
        new_string = new_string.replace(from, to);
    }
    new_string
}

/// unescapes all chars described in the server query manual
/// # Example
/// ```
/// use sqlib::escaping::unescape;
///
/// let unescaped = "hello world/|\\".to_string();
/// let escaped = "hello\\sworld\\/\\p\\\\".to_string();
///
/// let s = unescape(&escaped);
///
/// assert_eq!(s, unescaped);
/// ```
pub fn unescape(s: &str) -> String {
    let mut new_string = s.to_string();
    for &(to, from) in ESCAPE_CHARS.iter() {
        new_string = new_string.replace(from, &to.to_string());
    }
    new_string
}
<commit_msg>adding better documentation and an extra test to escaping.rs<commit_after>//! escaping provides functions for the Server Query escaping.
//!
//! # Example
//! ```
//! use sqlib::escaping::{escape, unescape};
//!
//! let unescaped = "hello world/|\\".to_string();
//! let escaped = "hello\\sworld\\/\\p\\\\".to_string();
//!
//! let escaped_test = escape(&unescaped);
//! let unescaped_test = unescape(&escaped);
//!
//! assert_eq!(escaped_test, escaped);
//! assert_eq!(unescaped_test, unescaped);
//! ```

const ESCAPE_CHARS: [(char, &'static str); 11] = [('\\', "\\\\"),
                                                  (' ', "\\s"),
                                                  ('/', "\\/"),
                                                  ('|', "\\p"),
                                                  (7 as char, "\\a"),
                                                  (8 as char, "\\b"),
                                                  (9 as char, "\\t"),
                                                  (10 as char, "\\n"),
                                                  (11 as char, "\\v"),
                                                  (12 as char, "\\f"),
                                                  (13 as char, "\\r")];

/// escapes all chars described in the server query manual
///
/// # Example
/// ```
/// use sqlib::escaping::escape;
///
/// let unescaped = "hello world/|\\".to_string();
/// let escaped = "hello\\sworld\\/\\p\\\\".to_string();
///
/// let s = escape(&unescaped);
///
/// assert_eq!(s, escaped);
/// ```
pub fn escape(s: &str) -> String {
    let mut new_string = s.to_string();
    for &(from, to) in ESCAPE_CHARS.iter() {
        new_string = new_string.replace(from, to);
    }
    new_string
}

/// unescapes all chars described in the server query manual
///
/// # Example
/// ```
/// use sqlib::escaping::unescape;
///
/// let unescaped = "hello world/|\\".to_string();
/// let escaped = "hello\\sworld\\/\\p\\\\".to_string();
///
/// let s = unescape(&escaped);
///
/// assert_eq!(s, unescaped);
/// ```
pub fn unescape(s: &str) -> String {
    let mut new_string = s.to_string();
    for &(to, from) in ESCAPE_CHARS.iter() {
        new_string = new_string.replace(from, &to.to_string());
    }
    new_string
}
<|endoftext|>


<commit_before>use crate::types::GossipKind;
use crate::{Enr, PeerIdSerialized};
use directory::{
    DEFAULT_BEACON_NODE_DIR, DEFAULT_HARDCODED_NETWORK, DEFAULT_NETWORK_DIR, DEFAULT_ROOT_DIR,
};
use discv5::{Discv5Config, Discv5ConfigBuilder};
use libp2p::gossipsub::{
    FastMessageId, GossipsubConfig, GossipsubConfigBuilder, GossipsubMessage, MessageId,
    RawGossipsubMessage, ValidationMode,
};
use libp2p::Multiaddr;
use serde_derive::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use types::{ForkContext, ForkName};

/// The maximum transmit size of gossip messages in bytes pre-merge.
const GOSSIP_MAX_SIZE: usize = 1_048_576; // 1M
/// The maximum transmit size of gossip messages in bytes post-merge.
const GOSSIP_MAX_SIZE_POST_MERGE: usize = 10 * 1_048_576; // 10M

/// The cache time is set to accommodate the circulation time of an attestation.
///
/// The p2p spec declares that we accept attestations within the following range:
///
/// ```ignore
/// ATTESTATION_PROPAGATION_SLOT_RANGE = 32
/// attestation.data.slot + ATTESTATION_PROPAGATION_SLOT_RANGE >= current_slot >= attestation.data.slot
/// ```
///
/// Therefore, we must accept attestations across a span of 33 slots (where each slot is 12
/// seconds). We add an additional second to account for the 500ms gossip clock disparity, and
/// another 500ms for "fudge factor".
pub const DUPLICATE_CACHE_TIME: Duration = Duration::from_secs(33 * 12 + 1);

// We treat uncompressed messages as invalid and never use the INVALID_SNAPPY_DOMAIN as in the
// specification. We leave it here for posterity.
// const MESSAGE_DOMAIN_INVALID_SNAPPY: [u8; 4] = [0, 0, 0, 0];
const MESSAGE_DOMAIN_VALID_SNAPPY: [u8; 4] = [1, 0, 0, 0];

/// The maximum size of gossip messages.
pub fn gossip_max_size(is_merge_enabled: bool) -> usize {
    if is_merge_enabled {
        GOSSIP_MAX_SIZE_POST_MERGE
    } else {
        GOSSIP_MAX_SIZE
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
/// Network configuration for lighthouse.
pub struct Config {
    /// Data directory where node's keyfile is stored
    pub network_dir: PathBuf,

    /// IP address to listen on.
    pub listen_address: std::net::IpAddr,

    /// The TCP port that libp2p listens on.
    pub libp2p_port: u16,

    /// UDP port that discovery listens on.
    pub discovery_port: u16,

    /// The address to broadcast to peers about which address we are listening on. None indicates
    /// that no discovery address has been set in the CLI args.
    pub enr_address: Option<std::net::IpAddr>,

    /// The udp port to broadcast to peers in order to reach back for discovery.
    pub enr_udp_port: Option<u16>,

    /// The tcp port to broadcast to peers in order to reach back for libp2p services.
    pub enr_tcp_port: Option<u16>,

    /// Target number of connected peers.
    pub target_peers: usize,

    /// Gossipsub configuration parameters.
    #[serde(skip)]
    pub gs_config: GossipsubConfig,

    /// Discv5 configuration parameters.
    #[serde(skip)]
    pub discv5_config: Discv5Config,

    /// List of nodes to initially connect to.
    pub boot_nodes_enr: Vec<Enr>,

    /// List of nodes to initially connect to, on Multiaddr format.
    pub boot_nodes_multiaddr: Vec<Multiaddr>,

    /// List of libp2p nodes to initially connect to.
    pub libp2p_nodes: Vec<Multiaddr>,

    /// List of trusted libp2p nodes which are not scored.
    pub trusted_peers: Vec<PeerIdSerialized>,

    /// Client version
    pub client_version: String,

    /// Disables the discovery protocol from starting.
    pub disable_discovery: bool,

    /// Attempt to construct external port mappings with UPnP.
    pub upnp_enabled: bool,

    /// Subscribe to all subnets for the duration of the runtime.
    pub subscribe_all_subnets: bool,

    /// Import/aggregate all attestations recieved on subscribed subnets for the duration of the
    /// runtime.
    pub import_all_attestations: bool,

    /// A setting specifying a range of values that tune the network parameters of lighthouse. The
    /// lower the value the less bandwidth used, but the slower messages will be received.
    pub network_load: u8,

    /// Indicates if the user has set the network to be in private mode. Currently this
    /// prevents sending client identifying information over identify.
    pub private: bool,

    /// Shutdown beacon node after sync is completed.
    pub shutdown_after_sync: bool,

    /// List of extra topics to initially subscribe to as strings.
    pub topics: Vec<GossipKind>,

    /// Whether metrics are enabled.
    pub metrics_enabled: bool,
}

impl Default for Config {
    /// Generate a default network configuration.
    fn default() -> Self {
        // WARNING: this directory default should be always overwritten with parameters
        // from cli for specific networks.
        let network_dir = dirs::home_dir()
            .unwrap_or_else(|| PathBuf::from("."))
            .join(DEFAULT_ROOT_DIR)
            .join(DEFAULT_HARDCODED_NETWORK)
            .join(DEFAULT_BEACON_NODE_DIR)
            .join(DEFAULT_NETWORK_DIR);

        // Note: Using the default config here. Use `gossipsub_config` function for getting
        // Lighthouse specific configuration for gossipsub.
        let gs_config = GossipsubConfigBuilder::default()
            .build()
            .expect("valid gossipsub configuration");

        // Discv5 Unsolicited Packet Rate Limiter
        let filter_rate_limiter = Some(
            discv5::RateLimiterBuilder::new()
                .total_n_every(10, Duration::from_secs(1)) // Allow bursts, average 10 per second
                .ip_n_every(9, Duration::from_secs(1)) // Allow bursts, average 9 per second
                .node_n_every(8, Duration::from_secs(1)) // Allow bursts, average 8 per second
                .build()
                .expect("The total rate limit has been specified"),
        );

        // discv5 configuration
        let discv5_config = Discv5ConfigBuilder::new()
            .enable_packet_filter()
            .session_cache_capacity(5000)
            .request_timeout(Duration::from_secs(1))
            .query_peer_timeout(Duration::from_secs(2))
            .query_timeout(Duration::from_secs(30))
            .request_retries(1)
            .enr_peer_update_min(10)
            .query_parallelism(5)
            .disable_report_discovered_peers()
            .ip_limit() // limits /24 IP's in buckets.
            .incoming_bucket_limit(8) // half the bucket size
            .filter_rate_limiter(filter_rate_limiter)
            .filter_max_bans_per_ip(Some(5))
            .filter_max_nodes_per_ip(Some(10))
            .ban_duration(Some(Duration::from_secs(3600)))
            .ping_interval(Duration::from_secs(300))
            .build();

        // NOTE: Some of these get overridden by the corresponding CLI default values.
        Config {
            network_dir,
            listen_address: "0.0.0.0".parse().expect("valid ip address"),
            libp2p_port: 9000,
            discovery_port: 9000,
            enr_address: None,
            enr_udp_port: None,
            enr_tcp_port: None,
            target_peers: 50,
            gs_config,
            discv5_config,
            boot_nodes_enr: vec![],
            boot_nodes_multiaddr: vec![],
            libp2p_nodes: vec![],
            trusted_peers: vec![],
            client_version: lighthouse_version::version_with_platform(),
            disable_discovery: false,
            upnp_enabled: true,
            network_load: 3,
            private: false,
            subscribe_all_subnets: false,
            import_all_attestations: false,
            shutdown_after_sync: false,
            topics: Vec::new(),
            metrics_enabled: false,
        }
    }
}

/// Controls sizes of gossipsub meshes to tune a Lighthouse node's bandwidth/performance.
pub struct NetworkLoad {
    pub name: &'static str,
    pub mesh_n_low: usize,
    pub outbound_min: usize,
    pub mesh_n: usize,
    pub mesh_n_high: usize,
    pub gossip_lazy: usize,
    pub history_gossip: usize,
}

impl From<u8> for NetworkLoad {
    fn from(load: u8) -> NetworkLoad {
        match load {
            1 => NetworkLoad {
                name: "Low",
                mesh_n_low: 1,
                outbound_min: 1,
                mesh_n: 3,
                mesh_n_high: 4,
                gossip_lazy: 3,
                history_gossip: 12,
            },
            2 => NetworkLoad {
                name: "Low",
                mesh_n_low: 2,
                outbound_min: 2,
                mesh_n: 4,
                mesh_n_high: 8,
                gossip_lazy: 3,
                history_gossip: 12,
            },
            3 => NetworkLoad {
                name: "Average",
                mesh_n_low: 3,
                outbound_min: 2,
                mesh_n: 5,
                mesh_n_high: 10,
                gossip_lazy: 3,
                history_gossip: 12,
            },
            4 => NetworkLoad {
                name: "Average",
                mesh_n_low: 4,
                outbound_min: 3,
                mesh_n: 8,
                mesh_n_high: 12,
                gossip_lazy: 3,
                history_gossip: 12,
            },
            // 5 and above
            _ => NetworkLoad {
                name: "High",
                mesh_n_low: 5,
                outbound_min: 3,
                mesh_n: 10,
                mesh_n_high: 15,
                gossip_lazy: 5,
                history_gossip: 12,
            },
        }
    }
}

/// Return a Lighthouse specific `GossipsubConfig` where the `message_id_fn` depends on the current fork.
pub fn gossipsub_config(network_load: u8, fork_context: Arc<ForkContext>) -> GossipsubConfig {
    // The function used to generate a gossipsub message id
    // We use the first 8 bytes of SHA256(data) for content addressing
    let fast_gossip_message_id =
        |message: &RawGossipsubMessage| FastMessageId::from(&Sha256::digest(&message.data)[..8]);
    fn prefix(
        prefix: [u8; 4],
        message: &GossipsubMessage,
        fork_context: Arc<ForkContext>,
    ) -> Vec<u8> {
        let topic_bytes = message.topic.as_str().as_bytes();
        match fork_context.current_fork() {
            // according to: https://github.com/ethereum/consensus-specs/blob/dev/specs/merge/p2p-interface.md#the-gossip-domain-gossipsub
            // the derivation of the message-id remains the same in the merge
            ForkName::Altair | ForkName::Merge => {
                let topic_len_bytes = topic_bytes.len().to_le_bytes();
                let mut vec = Vec::with_capacity(
                    prefix.len() + topic_len_bytes.len() + topic_bytes.len() + message.data.len(),
                );
                vec.extend_from_slice(&prefix);
                vec.extend_from_slice(&topic_len_bytes);
                vec.extend_from_slice(topic_bytes);
                vec.extend_from_slice(&message.data);
                vec
            }
            ForkName::Base => {
                let mut vec = Vec::with_capacity(prefix.len() + message.data.len());
                vec.extend_from_slice(&prefix);
                vec.extend_from_slice(&message.data);
                vec
            }
        }
    }

    let is_merge_enabled = fork_context.fork_exists(ForkName::Merge);
    let gossip_message_id = move |message: &GossipsubMessage| {
        MessageId::from(
            &Sha256::digest(
                prefix(MESSAGE_DOMAIN_VALID_SNAPPY, message, fork_context.clone()).as_slice(),
            )[..20],
        )
    };

    let load = NetworkLoad::from(network_load);

    GossipsubConfigBuilder::default()
        .max_transmit_size(gossip_max_size(is_merge_enabled))
        .heartbeat_interval(Duration::from_millis(700))
        .mesh_n(load.mesh_n)
        .mesh_n_low(load.mesh_n_low)
        .mesh_outbound_min(load.outbound_min)
        .mesh_n_high(load.mesh_n_high)
        .gossip_lazy(load.gossip_lazy)
        .fanout_ttl(Duration::from_secs(60))
        .history_length(12)
        .max_messages_per_rpc(Some(500)) // Responses to IWANT can be quite large
        .history_gossip(load.history_gossip)
        .validate_messages() // require validation before propagation
        .validation_mode(ValidationMode::Anonymous)
        .duplicate_cache_time(DUPLICATE_CACHE_TIME)
        .message_id_fn(gossip_message_id)
        .fast_message_id_fn(fast_gossip_message_id)
        .allow_self_origin(true)
        .build()
        .expect("valid gossipsub configuration")
}
<commit_msg>Reduce gossip history (#2969)<commit_after>use crate::types::GossipKind;
use crate::{Enr, PeerIdSerialized};
use directory::{
    DEFAULT_BEACON_NODE_DIR, DEFAULT_HARDCODED_NETWORK, DEFAULT_NETWORK_DIR, DEFAULT_ROOT_DIR,
};
use discv5::{Discv5Config, Discv5ConfigBuilder};
use libp2p::gossipsub::{
    FastMessageId, GossipsubConfig, GossipsubConfigBuilder, GossipsubMessage, MessageId,
    RawGossipsubMessage, ValidationMode,
};
use libp2p::Multiaddr;
use serde_derive::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use types::{ForkContext, ForkName};

/// The maximum transmit size of gossip messages in bytes pre-merge.
const GOSSIP_MAX_SIZE: usize = 1_048_576; // 1M
/// The maximum transmit size of gossip messages in bytes post-merge.
const GOSSIP_MAX_SIZE_POST_MERGE: usize = 10 * 1_048_576; // 10M

/// The cache time is set to accommodate the circulation time of an attestation.
///
/// The p2p spec declares that we accept attestations within the following range:
///
/// ```ignore
/// ATTESTATION_PROPAGATION_SLOT_RANGE = 32
/// attestation.data.slot + ATTESTATION_PROPAGATION_SLOT_RANGE >= current_slot >= attestation.data.slot
/// ```
///
/// Therefore, we must accept attestations across a span of 33 slots (where each slot is 12
/// seconds). We add an additional second to account for the 500ms gossip clock disparity, and
/// another 500ms for "fudge factor".
pub const DUPLICATE_CACHE_TIME: Duration = Duration::from_secs(33 * 12 + 1);

// We treat uncompressed messages as invalid and never use the INVALID_SNAPPY_DOMAIN as in the
// specification. We leave it here for posterity.
// const MESSAGE_DOMAIN_INVALID_SNAPPY: [u8; 4] = [0, 0, 0, 0];
const MESSAGE_DOMAIN_VALID_SNAPPY: [u8; 4] = [1, 0, 0, 0];

/// The maximum size of gossip messages.
pub fn gossip_max_size(is_merge_enabled: bool) -> usize {
    if is_merge_enabled {
        GOSSIP_MAX_SIZE_POST_MERGE
    } else {
        GOSSIP_MAX_SIZE
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
/// Network configuration for lighthouse.
pub struct Config {
    /// Data directory where node's keyfile is stored
    pub network_dir: PathBuf,

    /// IP address to listen on.
    pub listen_address: std::net::IpAddr,

    /// The TCP port that libp2p listens on.
    pub libp2p_port: u16,

    /// UDP port that discovery listens on.
    pub discovery_port: u16,

    /// The address to broadcast to peers about which address we are listening on. None indicates
    /// that no discovery address has been set in the CLI args.
    pub enr_address: Option<std::net::IpAddr>,

    /// The udp port to broadcast to peers in order to reach back for discovery.
    pub enr_udp_port: Option<u16>,

    /// The tcp port to broadcast to peers in order to reach back for libp2p services.
    pub enr_tcp_port: Option<u16>,

    /// Target number of connected peers.
    pub target_peers: usize,

    /// Gossipsub configuration parameters.
    #[serde(skip)]
    pub gs_config: GossipsubConfig,

    /// Discv5 configuration parameters.
    #[serde(skip)]
    pub discv5_config: Discv5Config,

    /// List of nodes to initially connect to.
    pub boot_nodes_enr: Vec<Enr>,

    /// List of nodes to initially connect to, on Multiaddr format.
    pub boot_nodes_multiaddr: Vec<Multiaddr>,

    /// List of libp2p nodes to initially connect to.
    pub libp2p_nodes: Vec<Multiaddr>,

    /// List of trusted libp2p nodes which are not scored.
    pub trusted_peers: Vec<PeerIdSerialized>,

    /// Client version
    pub client_version: String,

    /// Disables the discovery protocol from starting.
    pub disable_discovery: bool,

    /// Attempt to construct external port mappings with UPnP.
    pub upnp_enabled: bool,

    /// Subscribe to all subnets for the duration of the runtime.
    pub subscribe_all_subnets: bool,

    /// Import/aggregate all attestations recieved on subscribed subnets for the duration of the
    /// runtime.
    pub import_all_attestations: bool,

    /// A setting specifying a range of values that tune the network parameters of lighthouse. The
    /// lower the value the less bandwidth used, but the slower messages will be received.
    pub network_load: u8,

    /// Indicates if the user has set the network to be in private mode. Currently this
    /// prevents sending client identifying information over identify.
    pub private: bool,

    /// Shutdown beacon node after sync is completed.
    pub shutdown_after_sync: bool,

    /// List of extra topics to initially subscribe to as strings.
    pub topics: Vec<GossipKind>,

    /// Whether metrics are enabled.
    pub metrics_enabled: bool,
}

impl Default for Config {
    /// Generate a default network configuration.
    fn default() -> Self {
        // WARNING: this directory default should be always overwritten with parameters
        // from cli for specific networks.
        let network_dir = dirs::home_dir()
            .unwrap_or_else(|| PathBuf::from("."))
            .join(DEFAULT_ROOT_DIR)
            .join(DEFAULT_HARDCODED_NETWORK)
            .join(DEFAULT_BEACON_NODE_DIR)
            .join(DEFAULT_NETWORK_DIR);

        // Note: Using the default config here. Use `gossipsub_config` function for getting
        // Lighthouse specific configuration for gossipsub.
        let gs_config = GossipsubConfigBuilder::default()
            .build()
            .expect("valid gossipsub configuration");

        // Discv5 Unsolicited Packet Rate Limiter
        let filter_rate_limiter = Some(
            discv5::RateLimiterBuilder::new()
                .total_n_every(10, Duration::from_secs(1)) // Allow bursts, average 10 per second
                .ip_n_every(9, Duration::from_secs(1)) // Allow bursts, average 9 per second
                .node_n_every(8, Duration::from_secs(1)) // Allow bursts, average 8 per second
                .build()
                .expect("The total rate limit has been specified"),
        );

        // discv5 configuration
        let discv5_config = Discv5ConfigBuilder::new()
            .enable_packet_filter()
            .session_cache_capacity(5000)
            .request_timeout(Duration::from_secs(1))
            .query_peer_timeout(Duration::from_secs(2))
            .query_timeout(Duration::from_secs(30))
            .request_retries(1)
            .enr_peer_update_min(10)
            .query_parallelism(5)
            .disable_report_discovered_peers()
            .ip_limit() // limits /24 IP's in buckets.
            .incoming_bucket_limit(8) // half the bucket size
            .filter_rate_limiter(filter_rate_limiter)
            .filter_max_bans_per_ip(Some(5))
            .filter_max_nodes_per_ip(Some(10))
            .ban_duration(Some(Duration::from_secs(3600)))
            .ping_interval(Duration::from_secs(300))
            .build();

        // NOTE: Some of these get overridden by the corresponding CLI default values.
        Config {
            network_dir,
            listen_address: "0.0.0.0".parse().expect("valid ip address"),
            libp2p_port: 9000,
            discovery_port: 9000,
            enr_address: None,
            enr_udp_port: None,
            enr_tcp_port: None,
            target_peers: 50,
            gs_config,
            discv5_config,
            boot_nodes_enr: vec![],
            boot_nodes_multiaddr: vec![],
            libp2p_nodes: vec![],
            trusted_peers: vec![],
            client_version: lighthouse_version::version_with_platform(),
            disable_discovery: false,
            upnp_enabled: true,
            network_load: 3,
            private: false,
            subscribe_all_subnets: false,
            import_all_attestations: false,
            shutdown_after_sync: false,
            topics: Vec::new(),
            metrics_enabled: false,
        }
    }
}

/// Controls sizes of gossipsub meshes to tune a Lighthouse node's bandwidth/performance.
pub struct NetworkLoad {
    pub name: &'static str,
    pub mesh_n_low: usize,
    pub outbound_min: usize,
    pub mesh_n: usize,
    pub mesh_n_high: usize,
    pub gossip_lazy: usize,
    pub history_gossip: usize,
    pub heartbeat_interval: Duration,
}

impl From<u8> for NetworkLoad {
    fn from(load: u8) -> NetworkLoad {
        match load {
            1 => NetworkLoad {
                name: "Low",
                mesh_n_low: 1,
                outbound_min: 1,
                mesh_n: 3,
                mesh_n_high: 4,
                gossip_lazy: 3,
                history_gossip: 3,
                heartbeat_interval: Duration::from_millis(1200),
            },
            2 => NetworkLoad {
                name: "Low",
                mesh_n_low: 2,
                outbound_min: 2,
                mesh_n: 4,
                mesh_n_high: 8,
                gossip_lazy: 3,
                history_gossip: 3,
                heartbeat_interval: Duration::from_millis(1000),
            },
            3 => NetworkLoad {
                name: "Average",
                mesh_n_low: 3,
                outbound_min: 2,
                mesh_n: 5,
                mesh_n_high: 10,
                gossip_lazy: 3,
                history_gossip: 3,
                heartbeat_interval: Duration::from_millis(700),
            },
            4 => NetworkLoad {
                name: "Average",
                mesh_n_low: 4,
                outbound_min: 3,
                mesh_n: 8,
                mesh_n_high: 12,
                gossip_lazy: 3,
                history_gossip: 3,
                heartbeat_interval: Duration::from_millis(700),
            },
            // 5 and above
            _ => NetworkLoad {
                name: "High",
                mesh_n_low: 5,
                outbound_min: 3,
                mesh_n: 10,
                mesh_n_high: 15,
                gossip_lazy: 5,
                history_gossip: 6,
                heartbeat_interval: Duration::from_millis(500),
            },
        }
    }
}

/// Return a Lighthouse specific `GossipsubConfig` where the `message_id_fn` depends on the current fork.
pub fn gossipsub_config(network_load: u8, fork_context: Arc<ForkContext>) -> GossipsubConfig {
    // The function used to generate a gossipsub message id
    // We use the first 8 bytes of SHA256(data) for content addressing
    let fast_gossip_message_id =
        |message: &RawGossipsubMessage| FastMessageId::from(&Sha256::digest(&message.data)[..8]);
    fn prefix(
        prefix: [u8; 4],
        message: &GossipsubMessage,
        fork_context: Arc<ForkContext>,
    ) -> Vec<u8> {
        let topic_bytes = message.topic.as_str().as_bytes();
        match fork_context.current_fork() {
            // according to: https://github.com/ethereum/consensus-specs/blob/dev/specs/merge/p2p-interface.md#the-gossip-domain-gossipsub
            // the derivation of the message-id remains the same in the merge
            ForkName::Altair | ForkName::Merge => {
                let topic_len_bytes = topic_bytes.len().to_le_bytes();
                let mut vec = Vec::with_capacity(
                    prefix.len() + topic_len_bytes.len() + topic_bytes.len() + message.data.len(),
                );
                vec.extend_from_slice(&prefix);
                vec.extend_from_slice(&topic_len_bytes);
                vec.extend_from_slice(topic_bytes);
                vec.extend_from_slice(&message.data);
                vec
            }
            ForkName::Base => {
                let mut vec = Vec::with_capacity(prefix.len() + message.data.len());
                vec.extend_from_slice(&prefix);
                vec.extend_from_slice(&message.data);
                vec
            }
        }
    }

    let is_merge_enabled = fork_context.fork_exists(ForkName::Merge);
    let gossip_message_id = move |message: &GossipsubMessage| {
        MessageId::from(
            &Sha256::digest(
                prefix(MESSAGE_DOMAIN_VALID_SNAPPY, message, fork_context.clone()).as_slice(),
            )[..20],
        )
    };

    let load = NetworkLoad::from(network_load);

    GossipsubConfigBuilder::default()
        .max_transmit_size(gossip_max_size(is_merge_enabled))
        .heartbeat_interval(load.heartbeat_interval)
        .mesh_n(load.mesh_n)
        .mesh_n_low(load.mesh_n_low)
        .mesh_outbound_min(load.outbound_min)
        .mesh_n_high(load.mesh_n_high)
        .gossip_lazy(load.gossip_lazy)
        .fanout_ttl(Duration::from_secs(60))
        .history_length(12)
        .max_messages_per_rpc(Some(500)) // Responses to IWANT can be quite large
        .history_gossip(load.history_gossip)
        .validate_messages() // require validation before propagation
        .validation_mode(ValidationMode::Anonymous)
        .duplicate_cache_time(DUPLICATE_CACHE_TIME)
        .message_id_fn(gossip_message_id)
        .fast_message_id_fn(fast_gossip_message_id)
        .allow_self_origin(true)
        .build()
        .expect("valid gossipsub configuration")
}
<|endoftext|>