<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;

use proc_macro2::Span;
use syn::{
    parse::{Parse, ParseStream},
    parse_quote, LitStr, Token,
};
use syn::{Ident, Result as ParseResult};

#[derive(PartialEq, Clone, Debug, Hash)]
pub enum UnsafePolicy {
    AllFunctionsSafe,
    AllFunctionsUnsafe,
}

impl Parse for UnsafePolicy {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        if input.parse::<Option<Token![unsafe]>>()?.is_some() {
            return Ok(UnsafePolicy::AllFunctionsSafe);
        }
        let r = match input.parse::<Option<syn::Ident>>()? {
            Some(id) => {
                if id == "unsafe_ffi" {
                    Ok(UnsafePolicy::AllFunctionsSafe)
                } else {
                    Err(syn::Error::new(id.span(), "expected unsafe_ffi"))
                }
            }
            None => Ok(UnsafePolicy::AllFunctionsUnsafe),
        };
        if !input.is_empty() {
            return Err(syn::Error::new(
                Span::call_site(),
                "unexpected tokens within safety directive",
            ));
        }
        r
    }
}

/// Allowlist configuration.
#[derive(Hash, Debug)]
pub enum Allowlist {
    Unspecified(Vec<String>),
    All,
    Specific(Vec<String>),
}

impl Allowlist {
    pub(crate) fn push(&mut self, item: LitStr) -> ParseResult<()> {
        match self {
            Allowlist::Unspecified(ref mut uncommitted_list) => {
                let new_list = uncommitted_list
                    .drain(..)
                    .chain(std::iter::once(item.value()))
                    .collect();
                *self = Allowlist::Specific(new_list);
            }
            Allowlist::All => {
                return Err(syn::Error::new(
                    item.span(),
                    "use either generate!/generate_pod! or generate_all!, not both.",
                ))
            }
            Allowlist::Specific(list) => list.push(item.value()),
        };
        Ok(())
    }

    /// Sometimes we need to push things onto the allowlist without making
    /// a commitment whether we're in `allowlist` or `allowlist_all` mode.
    pub(crate) fn push_uncommitted(&mut self, item: LitStr) {
        match self {
            Allowlist::Unspecified(list) | Allowlist::Specific(list) => list.push(item.value()),
            Allowlist::All => {}
        }
    }

    pub(crate) fn set_all(&mut self, ident: &Ident) -> ParseResult<()> {
        if matches!(self, Allowlist::Specific(..)) {
            return Err(syn::Error::new(
                ident.span(),
                "use either generate!/generate_pod! or generate_all!, not both.",
            ));
        }
        *self = Allowlist::All;
        Ok(())
    }
}

impl Default for Allowlist {
    fn default() -> Self {
        Allowlist::Unspecified(Vec::new())
    }
}

#[derive(Hash, Debug)]
pub struct Subclass {
    pub superclass: String,
    pub subclass: Ident,
}

#[derive(Hash, Debug)]
pub struct IncludeCppConfig {
    pub inclusions: Vec<String>,
    pub unsafe_policy: UnsafePolicy,
    pub parse_only: bool,
    pub exclude_impls: bool,
    pod_requests: Vec<String>,
    allowlist: Allowlist,
    blocklist: Vec<String>,
    exclude_utilities: bool,
    mod_name: Option<Ident>,
    pub rust_types: Vec<Ident>,
    pub subclasses: Vec<Subclass>,
}

impl Parse for IncludeCppConfig {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut parse_only = false;
        let mut exclude_impls = false;
        let mut unsafe_policy = UnsafePolicy::AllFunctionsUnsafe;
        let mut allowlist = Allowlist::default();
        let mut blocklist = Vec::new();
        let mut pod_requests = Vec::new();
        let mut rust_types = Vec::new();
        let mut exclude_utilities = false;
        let mut mod_name = None;
        let mut subclasses = Vec::new();

        while !input.is_empty() {
            let has_hexathorpe = input.parse::<Option<syn::token::Pound>>()?.is_some();
            let ident: syn::Ident = input.parse()?;
            if has_hexathorpe {
                if ident != "include" {
                    return Err(syn::Error::new(ident.span(), "expected include"));
                }
                let hdr: syn::LitStr = input.parse()?;
                inclusions.push(hdr.value());
            } else {
                input.parse::<Option<syn::token::Bang>>()?;
                if ident == "generate" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    allowlist.push(generate)?;
                } else if ident == "generate_pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate_pod: syn::LitStr = args.parse()?;
                    pod_requests.push(generate_pod.value());
                    allowlist.push(generate_pod)?;
                } else if ident == "pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let pod: syn::LitStr = args.parse()?;
                    pod_requests.push(pod.value());
                } else if ident == "block" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    blocklist.push(generate.value());
                } else if ident == "rust_type" {
                    let args;
                    syn::parenthesized!(args in input);
                    let ident: syn::Ident = args.parse()?;
                    rust_types.push(ident);
                } else if ident == "subclass" {
                    let args;
                    syn::parenthesized!(args in input);
                    let superclass: syn::LitStr = args.parse()?;
                    args.parse::<syn::token::Comma>()?;
                    let subclass: syn::Ident = args.parse()?;
                    let sub_string = subclass.to_string();
                    subclasses.push(Subclass {
                        superclass: superclass.value(),
                        subclass,
                    });
                    allowlist.push_uncommitted(superclass);
                    allowlist.push_uncommitted(parse_quote! { #sub_string });
                } else if ident == "parse_only" {
                    parse_only = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "exclude_impls" {
                    exclude_impls = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "generate_all" {
                    allowlist.set_all(&ident)?;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "name" {
                    let args;
                    syn::parenthesized!(args in input);
                    let ident: syn::Ident = args.parse()?;
                    mod_name = Some(ident);
                } else if ident == "exclude_utilities" {
                    exclude_utilities = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "safety" {
                    let args;
                    syn::parenthesized!(args in input);
                    unsafe_policy = args.parse()?;
                } else {
                    return Err(syn::Error::new(
                        ident.span(),
                        "expected generate, generate_pod, nested_type, safety or exclude_utilities",
                    ));
                }
            }
            if input.is_empty() {
                break;
            }
        }

        Ok(IncludeCppConfig {
            inclusions,
            unsafe_policy,
            parse_only,
            exclude_impls,
            pod_requests,
            rust_types,
            allowlist,
            blocklist,
            exclude_utilities,
            mod_name,
            subclasses,
        })
    }
}

fn swallow_parentheses(input: &ParseStream, latest_ident: &Ident) -> ParseResult<()> {
    let args;
    syn::parenthesized!(args in input);
    if args.is_empty() {
        Ok(())
    } else {
        Err(syn::Error::new(
            latest_ident.span(),
            "expected no arguments to directive",
        ))
    }
}

impl IncludeCppConfig {
    pub fn get_pod_requests(&self) -> &[String] {
        &self.pod_requests
    }

    pub fn get_mod_name(&self) -> Ident {
        self.mod_name
            .as_ref()
            .cloned()
            .unwrap_or_else(|| Ident::new("ffi", Span::call_site()))
    }

    /// Whether to avoid generating the standard helpful utility
    /// functions which we normally include in every mod.
    pub fn exclude_utilities(&self) -> bool {
        self.exclude_utilities
    }

    /// Items which the user has explicitly asked us to generate;
    /// we should raise an error if we weren't able to do so.
    pub fn must_generate_list(&self) -> Box<dyn Iterator<Item = String> + '_> {
        if let Allowlist::Specific(items) = &self.allowlist {
            Box::new(items.iter().chain(self.pod_requests.iter()).cloned())
        } else {
            Box::new(self.pod_requests.iter().cloned())
        }
    }

    /// The allowlist of items to be passed into bindgen, if any.
    pub fn bindgen_allowlist(&self) -> Option<Box<dyn Iterator<Item = String> + '_>> {
        match &self.allowlist {
            Allowlist::All => None,
            Allowlist::Specific(items) => Some(Box::new(
                items
                    .iter()
                    .chain(self.pod_requests.iter())
                    .cloned()
                    .chain(self.active_utilities())
                    .chain(
                        self.subclasses
                            .iter()
                            .map(|sc| format!("{}Cpp", sc.subclass)),
                    ),
            )),
            Allowlist::Unspecified(_) => unreachable!(),
        }
    }

    fn active_utilities(&self) -> Vec<String> {
        if self.exclude_utilities {
            Vec::new()
        } else {
            vec![self.get_makestring_name()]
        }
    }

    /// Whether this type is on the allowlist specified by the user.
    ///
    /// A note on the allowlist handling in general. It's used in two places:
    /// 1) As directives to bindgen
    /// 2) After bindgen has generated code, to filter the APIs which
    ///    we pass to cxx.
    /// This second pass may seem redundant. But sometimes bindgen generates
    /// unnecessary stuff.
    pub fn is_on_allowlist(&self, cpp_name: &str) -> bool {
        match self.bindgen_allowlist() {
            None => true,
            Some(mut items) => {
                items.any(|item| item == cpp_name)
                    || self.active_utilities().iter().any(|item| *item == cpp_name)
            }
        }
    }

    pub fn is_on_blocklist(&self, cpp_name: &str) -> bool {
        self.blocklist.contains(&cpp_name.to_string())
    }

    pub fn get_blocklist(&self) -> impl Iterator<Item = &String> {
        self.blocklist.iter()
    }

    pub fn get_makestring_name(&self) -> String {
        format!(
            "autocxx_make_string_{}",
            self.mod_name
                .as_ref()
                .map(|i| i.to_string())
                .unwrap_or_else(|| "default".into())
        )
    }

    pub fn is_rust_type(&self, id: &Ident) -> bool {
        self.rust_types.contains(id)
    }

    pub fn superclasses(&self) -> impl Iterator<Item = &String> {
        let mut uniquified = HashSet::new();
        uniquified.extend(self.subclasses.iter().map(|sc| &sc.superclass));
        uniquified.into_iter()
    }
}

#[cfg(test)]
mod parse_tests {
    use crate::config::UnsafePolicy;
    use syn::parse_quote;
    #[test]
    fn test_safety_unsafe() {
        let us: UnsafePolicy = parse_quote! {
            unsafe
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_unsafe_ffi() {
        let us: UnsafePolicy = parse_quote! {
            unsafe_ffi
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_safe() {
        let us: UnsafePolicy = parse_quote! {};
        assert_eq!(us, UnsafePolicy::AllFunctionsUnsafe)
    }
}
<commit_msg>Simplify allowlist arrangement.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;

use proc_macro2::Span;
use syn::{
    parse::{Parse, ParseStream},
    LitStr, Token,
};
use syn::{Ident, Result as ParseResult};

#[derive(PartialEq, Clone, Debug, Hash)]
pub enum UnsafePolicy {
    AllFunctionsSafe,
    AllFunctionsUnsafe,
}

impl Parse for UnsafePolicy {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        if input.parse::<Option<Token![unsafe]>>()?.is_some() {
            return Ok(UnsafePolicy::AllFunctionsSafe);
        }
        let r = match input.parse::<Option<syn::Ident>>()? {
            Some(id) => {
                if id == "unsafe_ffi" {
                    Ok(UnsafePolicy::AllFunctionsSafe)
                } else {
                    Err(syn::Error::new(id.span(), "expected unsafe_ffi"))
                }
            }
            None => Ok(UnsafePolicy::AllFunctionsUnsafe),
        };
        if !input.is_empty() {
            return Err(syn::Error::new(
                Span::call_site(),
                "unexpected tokens within safety directive",
            ));
        }
        r
    }
}

/// Allowlist configuration.
#[derive(Hash, Debug)]
pub enum Allowlist {
    Unspecified(Vec<String>),
    All,
    Specific(Vec<String>),
}

impl Allowlist {
    pub(crate) fn push(&mut self, item: LitStr) -> ParseResult<()> {
        match self {
            Allowlist::Unspecified(ref mut uncommitted_list) => {
                let new_list = uncommitted_list
                    .drain(..)
                    .chain(std::iter::once(item.value()))
                    .collect();
                *self = Allowlist::Specific(new_list);
            }
            Allowlist::All => {
                return Err(syn::Error::new(
                    item.span(),
                    "use either generate!/generate_pod! or generate_all!, not both.",
                ))
            }
            Allowlist::Specific(list) => list.push(item.value()),
        };
        Ok(())
    }

    pub(crate) fn set_all(&mut self, ident: &Ident) -> ParseResult<()> {
        if matches!(self, Allowlist::Specific(..)) {
            return Err(syn::Error::new(
                ident.span(),
                "use either generate!/generate_pod! or generate_all!, not both.",
            ));
        }
        *self = Allowlist::All;
        Ok(())
    }
}

impl Default for Allowlist {
    fn default() -> Self {
        Allowlist::Unspecified(Vec::new())
    }
}

#[derive(Hash, Debug)]
pub struct Subclass {
    pub superclass: String,
    pub subclass: Ident,
}

#[derive(Hash, Debug)]
pub struct IncludeCppConfig {
    pub inclusions: Vec<String>,
    pub unsafe_policy: UnsafePolicy,
    pub parse_only: bool,
    pub exclude_impls: bool,
    pod_requests: Vec<String>,
    allowlist: Allowlist,
    blocklist: Vec<String>,
    exclude_utilities: bool,
    mod_name: Option<Ident>,
    pub rust_types: Vec<Ident>,
    pub subclasses: Vec<Subclass>,
}

impl Parse for IncludeCppConfig {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut parse_only = false;
        let mut exclude_impls = false;
        let mut unsafe_policy = UnsafePolicy::AllFunctionsUnsafe;
        let mut allowlist = Allowlist::default();
        let mut blocklist = Vec::new();
        let mut pod_requests = Vec::new();
        let mut rust_types = Vec::new();
        let mut exclude_utilities = false;
        let mut mod_name = None;
        let mut subclasses = Vec::new();

        while !input.is_empty() {
            let has_hexathorpe = input.parse::<Option<syn::token::Pound>>()?.is_some();
            let ident: syn::Ident = input.parse()?;
            if has_hexathorpe {
                if ident != "include" {
                    return Err(syn::Error::new(ident.span(), "expected include"));
                }
                let hdr: syn::LitStr = input.parse()?;
                inclusions.push(hdr.value());
            } else {
                input.parse::<Option<syn::token::Bang>>()?;
                if ident == "generate" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    allowlist.push(generate)?;
                } else if ident == "generate_pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate_pod: syn::LitStr = args.parse()?;
                    pod_requests.push(generate_pod.value());
                    allowlist.push(generate_pod)?;
                } else if ident == "pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let pod: syn::LitStr = args.parse()?;
                    pod_requests.push(pod.value());
                } else if ident == "block" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    blocklist.push(generate.value());
                } else if ident == "rust_type" {
                    let args;
                    syn::parenthesized!(args in input);
                    let ident: syn::Ident = args.parse()?;
                    rust_types.push(ident);
                } else if ident == "subclass" {
                    let args;
                    syn::parenthesized!(args in input);
                    let superclass: syn::LitStr = args.parse()?;
                    args.parse::<syn::token::Comma>()?;
                    let subclass: syn::Ident = args.parse()?;
                    subclasses.push(Subclass {
                        superclass: superclass.value(),
                        subclass,
                    });
                } else if ident == "parse_only" {
                    parse_only = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "exclude_impls" {
                    exclude_impls = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "generate_all" {
                    allowlist.set_all(&ident)?;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "name" {
                    let args;
                    syn::parenthesized!(args in input);
                    let ident: syn::Ident = args.parse()?;
                    mod_name = Some(ident);
                } else if ident == "exclude_utilities" {
                    exclude_utilities = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "safety" {
                    let args;
                    syn::parenthesized!(args in input);
                    unsafe_policy = args.parse()?;
                } else {
                    return Err(syn::Error::new(
                        ident.span(),
                        "expected generate, generate_pod, nested_type, safety or exclude_utilities",
                    ));
                }
            }
            if input.is_empty() {
                break;
            }
        }

        Ok(IncludeCppConfig {
            inclusions,
            unsafe_policy,
            parse_only,
            exclude_impls,
            pod_requests,
            rust_types,
            allowlist,
            blocklist,
            exclude_utilities,
            mod_name,
            subclasses,
        })
    }
}

fn swallow_parentheses(input: &ParseStream, latest_ident: &Ident) -> ParseResult<()> {
    let args;
    syn::parenthesized!(args in input);
    if args.is_empty() {
        Ok(())
    } else {
        Err(syn::Error::new(
            latest_ident.span(),
            "expected no arguments to directive",
        ))
    }
}

impl IncludeCppConfig {
    pub fn get_pod_requests(&self) -> &[String] {
        &self.pod_requests
    }

    pub fn get_mod_name(&self) -> Ident {
        self.mod_name
            .as_ref()
            .cloned()
            .unwrap_or_else(|| Ident::new("ffi", Span::call_site()))
    }

    /// Whether to avoid generating the standard helpful utility
    /// functions which we normally include in every mod.
    pub fn exclude_utilities(&self) -> bool {
        self.exclude_utilities
    }

    /// Items which the user has explicitly asked us to generate;
    /// we should raise an error if we weren't able to do so.
    pub fn must_generate_list(&self) -> Box<dyn Iterator<Item = String> + '_> {
        if let Allowlist::Specific(items) = &self.allowlist {
            Box::new(items.iter().chain(self.pod_requests.iter()).cloned())
        } else {
            Box::new(self.pod_requests.iter().cloned())
        }
    }

    /// The allowlist of items to be passed into bindgen, if any.
    pub fn bindgen_allowlist(&self) -> Option<Box<dyn Iterator<Item = String> + '_>> {
        match &self.allowlist {
            Allowlist::All => None,
            Allowlist::Specific(items) => Some(Box::new(
                items
                    .iter()
                    .chain(self.pod_requests.iter())
                    .cloned()
                    .chain(self.active_utilities())
                    .chain(
                        self.subclasses
                            .iter()
                            .map(|sc| {
                                [
                                    format!("{}Cpp", sc.subclass),
                                    sc.subclass.to_string(),
                                    sc.superclass.clone(),
                                ]
                            })
                            .flatten(),
                    ),
            )),
            Allowlist::Unspecified(_) => unreachable!(),
        }
    }

    fn active_utilities(&self) -> Vec<String> {
        if self.exclude_utilities {
            Vec::new()
        } else {
            vec![self.get_makestring_name()]
        }
    }

    /// Whether this type is on the allowlist specified by the user.
    ///
    /// A note on the allowlist handling in general. It's used in two places:
    /// 1) As directives to bindgen
    /// 2) After bindgen has generated code, to filter the APIs which
    ///    we pass to cxx.
    /// This second pass may seem redundant. But sometimes bindgen generates
    /// unnecessary stuff.
    pub fn is_on_allowlist(&self, cpp_name: &str) -> bool {
        match self.bindgen_allowlist() {
            None => true,
            Some(mut items) => {
                items.any(|item| item == cpp_name)
                    || self.active_utilities().iter().any(|item| *item == cpp_name)
            }
        }
    }

    pub fn is_on_blocklist(&self, cpp_name: &str) -> bool {
        self.blocklist.contains(&cpp_name.to_string())
    }

    pub fn get_blocklist(&self) -> impl Iterator<Item = &String> {
        self.blocklist.iter()
    }

    pub fn get_makestring_name(&self) -> String {
        format!(
            "autocxx_make_string_{}",
            self.mod_name
                .as_ref()
                .map(|i| i.to_string())
                .unwrap_or_else(|| "default".into())
        )
    }

    pub fn is_rust_type(&self, id: &Ident) -> bool {
        self.rust_types.contains(id)
    }

    pub fn superclasses(&self) -> impl Iterator<Item = &String> {
        let mut uniquified = HashSet::new();
        uniquified.extend(self.subclasses.iter().map(|sc| &sc.superclass));
        uniquified.into_iter()
    }
}

#[cfg(test)]
mod parse_tests {
    use crate::config::UnsafePolicy;
    use syn::parse_quote;
    #[test]
    fn test_safety_unsafe() {
        let us: UnsafePolicy = parse_quote! {
            unsafe
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_unsafe_ffi() {
        let us: UnsafePolicy = parse_quote! {
            unsafe_ffi
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_safe() {
        let us: UnsafePolicy = parse_quote! {};
        assert_eq!(us, UnsafePolicy::AllFunctionsUnsafe)
    }
}
<|endoftext|>


<commit_before>pub mod map;
mod navigate;
mod node;
mod search;
pub mod set;

#[doc(hidden)]
trait Recover<Q: ?Sized> {
    type Key;

    fn get(&self, key: &Q) -> Option<&Self::Key>;
    fn take(&mut self, key: &Q) -> Option<Self::Key>;
    fn replace(&mut self, key: Self::Key) -> Option<Self::Key>;
}

#[inline(always)]
pub unsafe fn unwrap_unchecked<T>(val: Option<T>) -> T {
    val.unwrap_or_else(|| {
        if cfg!(debug_assertions) {
            panic!("'unchecked' unwrap on None in BTreeMap");
        } else {
            unsafe {
                core::intrinsics::unreachable();
            }
        }
    })
}

#[cfg(test)]
/// XorShiftRng
struct DeterministicRng {
    x: u32,
    y: u32,
    z: u32,
    w: u32,
}

#[cfg(test)]
impl DeterministicRng {
    fn new() -> Self {
        DeterministicRng { x: 0x193a6754, y: 0xa8a7d469, z: 0x97830e05, w: 0x113ba7bb }
    }

    fn next(&mut self) -> u32 {
        let x = self.x;
        let t = x ^ (x << 11);
        self.x = self.y;
        self.y = self.z;
        self.z = self.w;
        let w_ = self.w;
        self.w = w_ ^ (w_ >> 19) ^ (t ^ (t >> 8));
        self.w
    }
}
<commit_msg>Rollup merge of #76543 - ssomers:btree_cleanup_4, r=Mark-Simulacrum<commit_after>pub mod map;
mod navigate;
mod node;
mod search;
pub mod set;

#[doc(hidden)]
trait Recover<Q: ?Sized> {
    type Key;

    fn get(&self, key: &Q) -> Option<&Self::Key>;
    fn take(&mut self, key: &Q) -> Option<Self::Key>;
    fn replace(&mut self, key: Self::Key) -> Option<Self::Key>;
}

/// Same purpose as `Option::unwrap` but doesn't always guarantee a panic
/// if the option contains no value.
/// SAFETY: the caller must ensure that the option contains a value.
#[inline(always)]
pub unsafe fn unwrap_unchecked<T>(val: Option<T>) -> T {
    val.unwrap_or_else(|| {
        if cfg!(debug_assertions) {
            panic!("'unchecked' unwrap on None in BTreeMap");
        } else {
            unsafe {
                core::intrinsics::unreachable();
            }
        }
    })
}

#[cfg(test)]
/// XorShiftRng
struct DeterministicRng {
    x: u32,
    y: u32,
    z: u32,
    w: u32,
}

#[cfg(test)]
impl DeterministicRng {
    fn new() -> Self {
        DeterministicRng { x: 0x193a6754, y: 0xa8a7d469, z: 0x97830e05, w: 0x113ba7bb }
    }

    fn next(&mut self) -> u32 {
        let x = self.x;
        let t = x ^ (x << 11);
        self.x = self.y;
        self.y = self.z;
        self.z = self.w;
        let w_ = self.w;
        self.w = w_ ^ (w_ >> 19) ^ (t ^ (t >> 8));
        self.w
    }
}
<|endoftext|>


<commit_before>/// This list represents the default file types that ripgrep ships with. In
/// general, any file format is fair game, although it should generally be
/// limited to reasonably popular open formats. For other cases, you can add
/// types to each invocation of ripgrep with the '--type-add' flag.
///
/// If you would like to add or improve this list, please file a PR:
/// https://github.com/BurntSushi/ripgrep
///
/// Please try to keep this list sorted lexicographically and wrapped to 79
/// columns (inclusive).
#[rustfmt::skip]
pub const DEFAULT_TYPES: &[(&str, &[&str])] = &[
    ("agda", &["*.agda", "*.lagda"]),
    ("aidl", &["*.aidl"]),
    ("amake", &["*.mk", "*.bp"]),
    ("asciidoc", &["*.adoc", "*.asc", "*.asciidoc"]),
    ("asm", &["*.asm", "*.s", "*.S"]),
    ("asp", &[
        "*.aspx", "*.aspx.cs", "*.aspx.cs", "*.ascx", "*.ascx.cs", "*.ascx.vb",
    ]),
    ("ats", &["*.ats", "*.dats", "*.sats", "*.hats"]),
    ("avro", &["*.avdl", "*.avpr", "*.avsc"]),
    ("awk", &["*.awk"]),
    ("bazel", &["*.bzl", "WORKSPACE", "BUILD", "BUILD.bazel"]),
    ("bitbake", &["*.bb", "*.bbappend", "*.bbclass", "*.conf", "*.inc"]),
    ("brotli", &["*.br"]),
    ("buildstream", &["*.bst"]),
    ("bzip2", &["*.bz2", "*.tbz2"]),
    ("c", &["*.[chH]", "*.[chH].in", "*.cats"]),
    ("cabal", &["*.cabal"]),
    ("cbor", &["*.cbor"]),
    ("ceylon", &["*.ceylon"]),
    ("clojure", &["*.clj", "*.cljc", "*.cljs", "*.cljx"]),
    ("cmake", &["*.cmake", "CMakeLists.txt"]),
    ("coffeescript", &["*.coffee"]),
    ("config", &["*.cfg", "*.conf", "*.config", "*.ini"]),
    ("cpp", &[
        "*.[ChH]", "*.cc", "*.[ch]pp", "*.[ch]xx", "*.hh",  "*.inl",
        "*.[ChH].in", "*.cc.in", "*.[ch]pp.in", "*.[ch]xx.in", "*.hh.in",
    ]),
    ("creole", &["*.creole"]),
    ("crystal", &["Projectfile", "*.cr"]),
    ("cs", &["*.cs"]),
    ("csharp", &["*.cs"]),
    ("cshtml", &["*.cshtml"]),
    ("css", &["*.css", "*.scss"]),
    ("csv", &["*.csv"]),
    ("cython", &["*.pyx", "*.pxi", "*.pxd"]),
    ("d", &["*.d"]),
    ("dart", &["*.dart"]),
    ("dhall", &["*.dhall"]),
    ("diff", &["*.patch", "*.diff"]),
    ("docker", &["*Dockerfile*"]),
    ("edn", &["*.edn"]),
    ("elisp", &["*.el"]),
    ("elixir", &["*.ex", "*.eex", "*.exs"]),
    ("elm", &["*.elm"]),
    ("erb", &["*.erb"]),
    ("erlang", &["*.erl", "*.hrl"]),
    ("fidl", &["*.fidl"]),
    ("fish", &["*.fish"]),
    ("fortran", &[
        "*.f", "*.F", "*.f77", "*.F77", "*.pfo",
        "*.f90", "*.F90", "*.f95", "*.F95",
    ]),
    ("fsharp", &["*.fs", "*.fsx", "*.fsi"]),
    ("gap", &["*.g", "*.gap", "*.gi", "*.gd", "*.tst"]),
    ("gn", &["*.gn", "*.gni"]),
    ("go", &["*.go"]),
    ("gradle", &["*.gradle"]),
    ("groovy", &["*.groovy", "*.gradle"]),
    ("gzip", &["*.gz", "*.tgz"]),
    ("h", &["*.h", "*.hpp"]),
    ("haml", &["*.haml"]),
    ("haskell", &["*.hs", "*.lhs", "*.cpphs", "*.c2hs", "*.hsc"]),
    ("hbs", &["*.hbs"]),
    ("hs", &["*.hs", "*.lhs"]),
    ("html", &["*.htm", "*.html", "*.ejs"]),
    ("idris", &["*.idr", "*.lidr"]),
    ("java", &["*.java", "*.jsp", "*.jspx", "*.properties"]),
    ("jinja", &["*.j2", "*.jinja", "*.jinja2"]),
    ("jl", &["*.jl"]),
    ("js", &["*.js", "*.jsx", "*.vue"]),
    ("json", &["*.json", "composer.lock"]),
    ("jsonl", &["*.jsonl"]),
    ("julia", &["*.jl"]),
    ("jupyter", &["*.ipynb", "*.jpynb"]),
    ("k", &["*.k"]),
    ("kotlin", &["*.kt", "*.kts"]),
    ("less", &["*.less"]),
    ("license", &[
        // General
        "COPYING", "COPYING[.-]*",
        "COPYRIGHT", "COPYRIGHT[.-]*",
        "EULA", "EULA[.-]*",
        "licen[cs]e", "licen[cs]e.*",
        "LICEN[CS]E", "LICEN[CS]E[.-]*", "*[.-]LICEN[CS]E*",
        "NOTICE", "NOTICE[.-]*",
        "PATENTS", "PATENTS[.-]*",
        "UNLICEN[CS]E", "UNLICEN[CS]E[.-]*",
        // GPL (gpl.txt, etc.)
        "agpl[.-]*",
        "gpl[.-]*",
        "lgpl[.-]*",
        // Other license-specific (APACHE-2.0.txt, etc.)
        "AGPL-*[0-9]*",
        "APACHE-*[0-9]*",
        "BSD-*[0-9]*",
        "CC-BY-*",
        "GFDL-*[0-9]*",
        "GNU-*[0-9]*",
        "GPL-*[0-9]*",
        "LGPL-*[0-9]*",
        "MIT-*[0-9]*",
        "MPL-*[0-9]*",
        "OFL-*[0-9]*",
    ]),
    ("lisp", &["*.el", "*.jl", "*.lisp", "*.lsp", "*.sc", "*.scm"]),
    ("lock", &["*.lock", "package-lock.json"]),
    ("log", &["*.log"]),
    ("lua", &["*.lua"]),
    ("lz4", &["*.lz4"]),
    ("lzma", &["*.lzma"]),
    ("m4", &["*.ac", "*.m4"]),
    ("make", &[
        "[Gg][Nn][Uu]makefile", "[Mm]akefile",
        "[Gg][Nn][Uu]makefile.am", "[Mm]akefile.am",
        "[Gg][Nn][Uu]makefile.in", "[Mm]akefile.in",
        "*.mk", "*.mak"
    ]),
    ("mako", &["*.mako", "*.mao"]),
    ("man", &["*.[0-9lnpx]", "*.[0-9][cEFMmpSx]"]),
    ("markdown", &["*.markdown", "*.md", "*.mdown", "*.mkdn"]),
    ("matlab", &["*.m"]),
    ("md", &["*.markdown", "*.md", "*.mdown", "*.mkdn"]),
    ("mk", &["mkfile"]),
    ("ml", &["*.ml"]),
    ("msbuild", &[
        "*.csproj", "*.fsproj", "*.vcxproj", "*.proj", "*.props", "*.targets",
    ]),
    ("nim", &["*.nim", "*.nimf", "*.nimble", "*.nims"]),
    ("nix", &["*.nix"]),
    ("objc", &["*.h", "*.m"]),
    ("objcpp", &["*.h", "*.mm"]),
    ("ocaml", &["*.ml", "*.mli", "*.mll", "*.mly"]),
    ("org", &["*.org", "*.org_archive"]),
    ("pascal", &["*.pas", "*.dpr", "*.lpr", "*.pp", "*.inc"]),
    ("pdf", &["*.pdf"]),
    ("perl", &["*.perl", "*.pl", "*.PL", "*.plh", "*.plx", "*.pm", "*.t"]),
    ("php", &["*.php", "*.php3", "*.php4", "*.php5", "*.phtml"]),
    ("pod", &["*.pod"]),
    ("postscript", &["*.eps", "*.ps"]),
    ("protobuf", &["*.proto"]),
    ("ps", &["*.cdxml", "*.ps1", "*.ps1xml", "*.psd1", "*.psm1"]),
    ("puppet", &["*.erb", "*.pp", "*.rb"]),
    ("purs", &["*.purs"]),
    ("py", &["*.py"]),
    ("qmake", &["*.pro", "*.pri", "*.prf"]),
    ("qml", &["*.qml"]),
    ("r", &["*.R", "*.r", "*.Rmd", "*.Rnw"]),
    ("rdoc", &["*.rdoc"]),
    ("readme", &["README*", "*README"]),
    ("robot", &["*.robot"]),
    ("rst", &["*.rst"]),
    ("ruby", &["Gemfile", "*.gemspec", ".irbrc", "Rakefile", "*.rb"]),
    ("rust", &["*.rs"]),
    ("sass", &["*.sass", "*.scss"]),
    ("scala", &["*.scala", "*.sbt"]),
    ("sh", &[
        // Portable/misc. init files
        ".login", ".logout", ".profile", "profile",
        // bash-specific init files
        ".bash_login", "bash_login",
        ".bash_logout", "bash_logout",
        ".bash_profile", "bash_profile",
        ".bashrc", "bashrc", "*.bashrc",
        // csh-specific init files
        ".cshrc", "*.cshrc",
        // ksh-specific init files
        ".kshrc", "*.kshrc",
        // tcsh-specific init files
        ".tcshrc",
        // zsh-specific init files
        ".zshenv", "zshenv",
        ".zlogin", "zlogin",
        ".zlogout", "zlogout",
        ".zprofile", "zprofile",
        ".zshrc", "zshrc",
        // Extensions
        "*.bash", "*.csh", "*.ksh", "*.sh", "*.tcsh", "*.zsh",
    ]),
    ("slim", &["*.skim", "*.slim", "*.slime"]),
    ("smarty", &["*.tpl"]),
    ("sml", &["*.sml", "*.sig"]),
    ("soy", &["*.soy"]),
    ("spark", &["*.spark"]),
    ("spec", &["*.spec"]),
    ("sql", &["*.sql", "*.psql"]),
    ("stylus", &["*.styl"]),
    ("sv", &["*.v", "*.vg", "*.sv", "*.svh", "*.h"]),
    ("svg", &["*.svg"]),
    ("swift", &["*.swift"]),
    ("swig", &["*.def", "*.i"]),
    ("systemd", &[
        "*.automount", "*.conf", "*.device", "*.link", "*.mount", "*.path",
        "*.scope", "*.service", "*.slice", "*.socket", "*.swap", "*.target",
        "*.timer",
    ]),
    ("taskpaper", &["*.taskpaper"]),
    ("tcl", &["*.tcl"]),
    ("tex", &["*.tex", "*.ltx", "*.cls", "*.sty", "*.bib", "*.dtx", "*.ins"]),
    ("textile", &["*.textile"]),
    ("tf", &["*.tf"]),
    ("thrift", &["*.thrift"]),
    ("toml", &["*.toml", "Cargo.lock"]),
    ("ts", &["*.ts", "*.tsx"]),
    ("twig", &["*.twig"]),
    ("txt", &["*.txt"]),
    ("typoscript", &["*.typoscript", "*.ts"]),
    ("vala", &["*.vala"]),
    ("vb", &["*.vb"]),
    ("verilog", &["*.v", "*.vh", "*.sv", "*.svh"]),
    ("vhdl", &["*.vhd", "*.vhdl"]),
    ("vim", &["*.vim"]),
    ("vimscript", &["*.vim"]),
    ("webidl", &["*.idl", "*.webidl", "*.widl"]),
    ("wiki", &["*.mediawiki", "*.wiki"]),
    ("xml", &[
        "*.xml", "*.xml.dist", "*.dtd", "*.xsl", "*.xslt", "*.xsd", "*.xjb",
        "*.rng", "*.sch", "*.xhtml",
    ]),
    ("xz", &["*.xz", "*.txz"]),
    ("yacc", &["*.y"]),
    ("yaml", &["*.yaml", "*.yml"]),
    ("zig", &["*.zig"]),
    ("zsh", &[
        ".zshenv", "zshenv",
        ".zlogin", "zlogin",
        ".zlogout", "zlogout",
        ".zprofile", "zprofile",
        ".zshrc", "zshrc",
        "*.zsh",
    ]),
    ("zstd", &["*.zst", "*.zstd"]),
];
<commit_msg>ignore/types: add coq type<commit_after>/// This list represents the default file types that ripgrep ships with. In
/// general, any file format is fair game, although it should generally be
/// limited to reasonably popular open formats. For other cases, you can add
/// types to each invocation of ripgrep with the '--type-add' flag.
///
/// If you would like to add or improve this list, please file a PR:
/// https://github.com/BurntSushi/ripgrep
///
/// Please try to keep this list sorted lexicographically and wrapped to 79
/// columns (inclusive).
#[rustfmt::skip]
pub const DEFAULT_TYPES: &[(&str, &[&str])] = &[
    ("agda", &["*.agda", "*.lagda"]),
    ("aidl", &["*.aidl"]),
    ("amake", &["*.mk", "*.bp"]),
    ("asciidoc", &["*.adoc", "*.asc", "*.asciidoc"]),
    ("asm", &["*.asm", "*.s", "*.S"]),
    ("asp", &[
        "*.aspx", "*.aspx.cs", "*.aspx.cs", "*.ascx", "*.ascx.cs", "*.ascx.vb",
    ]),
    ("ats", &["*.ats", "*.dats", "*.sats", "*.hats"]),
    ("avro", &["*.avdl", "*.avpr", "*.avsc"]),
    ("awk", &["*.awk"]),
    ("bazel", &["*.bzl", "WORKSPACE", "BUILD", "BUILD.bazel"]),
    ("bitbake", &["*.bb", "*.bbappend", "*.bbclass", "*.conf", "*.inc"]),
    ("brotli", &["*.br"]),
    ("buildstream", &["*.bst"]),
    ("bzip2", &["*.bz2", "*.tbz2"]),
    ("c", &["*.[chH]", "*.[chH].in", "*.cats"]),
    ("cabal", &["*.cabal"]),
    ("cbor", &["*.cbor"]),
    ("ceylon", &["*.ceylon"]),
    ("clojure", &["*.clj", "*.cljc", "*.cljs", "*.cljx"]),
    ("cmake", &["*.cmake", "CMakeLists.txt"]),
    ("coffeescript", &["*.coffee"]),
    ("config", &["*.cfg", "*.conf", "*.config", "*.ini"]),
    ("coq", &["*.v"]),
    ("cpp", &[
        "*.[ChH]", "*.cc", "*.[ch]pp", "*.[ch]xx", "*.hh",  "*.inl",
        "*.[ChH].in", "*.cc.in", "*.[ch]pp.in", "*.[ch]xx.in", "*.hh.in",
    ]),
    ("creole", &["*.creole"]),
    ("crystal", &["Projectfile", "*.cr"]),
    ("cs", &["*.cs"]),
    ("csharp", &["*.cs"]),
    ("cshtml", &["*.cshtml"]),
    ("css", &["*.css", "*.scss"]),
    ("csv", &["*.csv"]),
    ("cython", &["*.pyx", "*.pxi", "*.pxd"]),
    ("d", &["*.d"]),
    ("dart", &["*.dart"]),
    ("dhall", &["*.dhall"]),
    ("diff", &["*.patch", "*.diff"]),
    ("docker", &["*Dockerfile*"]),
    ("edn", &["*.edn"]),
    ("elisp", &["*.el"]),
    ("elixir", &["*.ex", "*.eex", "*.exs"]),
    ("elm", &["*.elm"]),
    ("erb", &["*.erb"]),
    ("erlang", &["*.erl", "*.hrl"]),
    ("fidl", &["*.fidl"]),
    ("fish", &["*.fish"]),
    ("fortran", &[
        "*.f", "*.F", "*.f77", "*.F77", "*.pfo",
        "*.f90", "*.F90", "*.f95", "*.F95",
    ]),
    ("fsharp", &["*.fs", "*.fsx", "*.fsi"]),
    ("gap", &["*.g", "*.gap", "*.gi", "*.gd", "*.tst"]),
    ("gn", &["*.gn", "*.gni"]),
    ("go", &["*.go"]),
    ("gradle", &["*.gradle"]),
    ("groovy", &["*.groovy", "*.gradle"]),
    ("gzip", &["*.gz", "*.tgz"]),
    ("h", &["*.h", "*.hpp"]),
    ("haml", &["*.haml"]),
    ("haskell", &["*.hs", "*.lhs", "*.cpphs", "*.c2hs", "*.hsc"]),
    ("hbs", &["*.hbs"]),
    ("hs", &["*.hs", "*.lhs"]),
    ("html", &["*.htm", "*.html", "*.ejs"]),
    ("idris", &["*.idr", "*.lidr"]),
    ("java", &["*.java", "*.jsp", "*.jspx", "*.properties"]),
    ("jinja", &["*.j2", "*.jinja", "*.jinja2"]),
    ("jl", &["*.jl"]),
    ("js", &["*.js", "*.jsx", "*.vue"]),
    ("json", &["*.json", "composer.lock"]),
    ("jsonl", &["*.jsonl"]),
    ("julia", &["*.jl"]),
    ("jupyter", &["*.ipynb", "*.jpynb"]),
    ("k", &["*.k"]),
    ("kotlin", &["*.kt", "*.kts"]),
    ("less", &["*.less"]),
    ("license", &[
        // General
        "COPYING", "COPYING[.-]*",
        "COPYRIGHT", "COPYRIGHT[.-]*",
        "EULA", "EULA[.-]*",
        "licen[cs]e", "licen[cs]e.*",
        "LICEN[CS]E", "LICEN[CS]E[.-]*", "*[.-]LICEN[CS]E*",
        "NOTICE", "NOTICE[.-]*",
        "PATENTS", "PATENTS[.-]*",
        "UNLICEN[CS]E", "UNLICEN[CS]E[.-]*",
        // GPL (gpl.txt, etc.)
        "agpl[.-]*",
        "gpl[.-]*",
        "lgpl[.-]*",
        // Other license-specific (APACHE-2.0.txt, etc.)
        "AGPL-*[0-9]*",
        "APACHE-*[0-9]*",
        "BSD-*[0-9]*",
        "CC-BY-*",
        "GFDL-*[0-9]*",
        "GNU-*[0-9]*",
        "GPL-*[0-9]*",
        "LGPL-*[0-9]*",
        "MIT-*[0-9]*",
        "MPL-*[0-9]*",
        "OFL-*[0-9]*",
    ]),
    ("lisp", &["*.el", "*.jl", "*.lisp", "*.lsp", "*.sc", "*.scm"]),
    ("lock", &["*.lock", "package-lock.json"]),
    ("log", &["*.log"]),
    ("lua", &["*.lua"]),
    ("lz4", &["*.lz4"]),
    ("lzma", &["*.lzma"]),
    ("m4", &["*.ac", "*.m4"]),
    ("make", &[
        "[Gg][Nn][Uu]makefile", "[Mm]akefile",
        "[Gg][Nn][Uu]makefile.am", "[Mm]akefile.am",
        "[Gg][Nn][Uu]makefile.in", "[Mm]akefile.in",
        "*.mk", "*.mak"
    ]),
    ("mako", &["*.mako", "*.mao"]),
    ("man", &["*.[0-9lnpx]", "*.[0-9][cEFMmpSx]"]),
    ("markdown", &["*.markdown", "*.md", "*.mdown", "*.mkdn"]),
    ("matlab", &["*.m"]),
    ("md", &["*.markdown", "*.md", "*.mdown", "*.mkdn"]),
    ("mk", &["mkfile"]),
    ("ml", &["*.ml"]),
    ("msbuild", &[
        "*.csproj", "*.fsproj", "*.vcxproj", "*.proj", "*.props", "*.targets",
    ]),
    ("nim", &["*.nim", "*.nimf", "*.nimble", "*.nims"]),
    ("nix", &["*.nix"]),
    ("objc", &["*.h", "*.m"]),
    ("objcpp", &["*.h", "*.mm"]),
    ("ocaml", &["*.ml", "*.mli", "*.mll", "*.mly"]),
    ("org", &["*.org", "*.org_archive"]),
    ("pascal", &["*.pas", "*.dpr", "*.lpr", "*.pp", "*.inc"]),
    ("pdf", &["*.pdf"]),
    ("perl", &["*.perl", "*.pl", "*.PL", "*.plh", "*.plx", "*.pm", "*.t"]),
    ("php", &["*.php", "*.php3", "*.php4", "*.php5", "*.phtml"]),
    ("pod", &["*.pod"]),
    ("postscript", &["*.eps", "*.ps"]),
    ("protobuf", &["*.proto"]),
    ("ps", &["*.cdxml", "*.ps1", "*.ps1xml", "*.psd1", "*.psm1"]),
    ("puppet", &["*.erb", "*.pp", "*.rb"]),
    ("purs", &["*.purs"]),
    ("py", &["*.py"]),
    ("qmake", &["*.pro", "*.pri", "*.prf"]),
    ("qml", &["*.qml"]),
    ("r", &["*.R", "*.r", "*.Rmd", "*.Rnw"]),
    ("rdoc", &["*.rdoc"]),
    ("readme", &["README*", "*README"]),
    ("robot", &["*.robot"]),
    ("rst", &["*.rst"]),
    ("ruby", &["Gemfile", "*.gemspec", ".irbrc", "Rakefile", "*.rb"]),
    ("rust", &["*.rs"]),
    ("sass", &["*.sass", "*.scss"]),
    ("scala", &["*.scala", "*.sbt"]),
    ("sh", &[
        // Portable/misc. init files
        ".login", ".logout", ".profile", "profile",
        // bash-specific init files
        ".bash_login", "bash_login",
        ".bash_logout", "bash_logout",
        ".bash_profile", "bash_profile",
        ".bashrc", "bashrc", "*.bashrc",
        // csh-specific init files
        ".cshrc", "*.cshrc",
        // ksh-specific init files
        ".kshrc", "*.kshrc",
        // tcsh-specific init files
        ".tcshrc",
        // zsh-specific init files
        ".zshenv", "zshenv",
        ".zlogin", "zlogin",
        ".zlogout", "zlogout",
        ".zprofile", "zprofile",
        ".zshrc", "zshrc",
        // Extensions
        "*.bash", "*.csh", "*.ksh", "*.sh", "*.tcsh", "*.zsh",
    ]),
    ("slim", &["*.skim", "*.slim", "*.slime"]),
    ("smarty", &["*.tpl"]),
    ("sml", &["*.sml", "*.sig"]),
    ("soy", &["*.soy"]),
    ("spark", &["*.spark"]),
    ("spec", &["*.spec"]),
    ("sql", &["*.sql", "*.psql"]),
    ("stylus", &["*.styl"]),
    ("sv", &["*.v", "*.vg", "*.sv", "*.svh", "*.h"]),
    ("svg", &["*.svg"]),
    ("swift", &["*.swift"]),
    ("swig", &["*.def", "*.i"]),
    ("systemd", &[
        "*.automount", "*.conf", "*.device", "*.link", "*.mount", "*.path",
        "*.scope", "*.service", "*.slice", "*.socket", "*.swap", "*.target",
        "*.timer",
    ]),
    ("taskpaper", &["*.taskpaper"]),
    ("tcl", &["*.tcl"]),
    ("tex", &["*.tex", "*.ltx", "*.cls", "*.sty", "*.bib", "*.dtx", "*.ins"]),
    ("textile", &["*.textile"]),
    ("tf", &["*.tf"]),
    ("thrift", &["*.thrift"]),
    ("toml", &["*.toml", "Cargo.lock"]),
    ("ts", &["*.ts", "*.tsx"]),
    ("twig", &["*.twig"]),
    ("txt", &["*.txt"]),
    ("typoscript", &["*.typoscript", "*.ts"]),
    ("vala", &["*.vala"]),
    ("vb", &["*.vb"]),
    ("verilog", &["*.v", "*.vh", "*.sv", "*.svh"]),
    ("vhdl", &["*.vhd", "*.vhdl"]),
    ("vim", &["*.vim"]),
    ("vimscript", &["*.vim"]),
    ("webidl", &["*.idl", "*.webidl", "*.widl"]),
    ("wiki", &["*.mediawiki", "*.wiki"]),
    ("xml", &[
        "*.xml", "*.xml.dist", "*.dtd", "*.xsl", "*.xslt", "*.xsd", "*.xjb",
        "*.rng", "*.sch", "*.xhtml",
    ]),
    ("xz", &["*.xz", "*.txz"]),
    ("yacc", &["*.y"]),
    ("yaml", &["*.yaml", "*.yml"]),
    ("zig", &["*.zig"]),
    ("zsh", &[
        ".zshenv", "zshenv",
        ".zlogin", "zlogin",
        ".zlogout", "zlogout",
        ".zprofile", "zprofile",
        ".zshrc", "zshrc",
        "*.zsh",
    ]),
    ("zstd", &["*.zst", "*.zstd"]),
];
<|endoftext|>


<commit_before>use actix::prelude::*;
use godcoin::prelude::*;

mod common;
pub use common::*;

#[test]
fn transfer_from_minter() {
    System::run(|| {
        let minter = TestMinter::new();
        let from_addr = ScriptHash::from(&minter.genesis_info().script);
        let to_addr = KeyPair::gen();

        let create_tx = |fee: &str, amount: Asset| {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, fee),
                from: from_addr.clone(),
                to: (&to_addr.0).into(),
                amount,
                memo: vec![],
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };

        let from_bal = minter.chain().get_balance(&from_addr, &[]).unwrap();
        let amount = get_asset("1.0000 GRAEL");
        let tx = create_tx("1.0000 GRAEL", amount);
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.and_then(move |res| {
                assert_eq!(res, MsgResponse::Broadcast());
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let chain = minter.chain();
                let cur_bal = chain.get_balance(&to_addr.0.into(), &[]);
                assert_eq!(cur_bal, Some(amount));

                // The fee transfers back to the minter wallet in the form of a reward tx so it
                // must not be subtracted during the assertion
                let cur_bal = chain.get_balance(&from_addr, &[]);
                assert_eq!(cur_bal, from_bal.sub(amount));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn transfer_from_user() {
    System::run(|| {
        let minter = TestMinter::new();

        let user_1_addr = KeyPair::gen();
        let user_2_addr = KeyPair::gen();

        let fut = {
            let tx = {
                let mut tx = TransferTx {
                    base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                    from: ScriptHash::from(&minter.genesis_info().script),
                    to: (&user_1_addr.0).into(),
                    amount: get_asset("100.0000 GRAEL"),
                    memo: vec![],
                    script: minter.genesis_info().script.clone(),
                };
                tx.append_sign(&minter.genesis_info().wallet_keys[3]);
                tx.append_sign(&minter.genesis_info().wallet_keys[0]);
                TxVariant::TransferTx(tx)
            };
            minter.request(MsgRequest::Broadcast(tx))
        };

        System::current().arbiter().send(
            fut.and_then({
                let user_1_addr = user_1_addr.clone();
                let user_2_addr = user_2_addr.clone();
                move |res| {
                    assert_eq!(res, MsgResponse::Broadcast());
                    let tx = {
                        let mut tx = TransferTx {
                            base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                            from: (&user_1_addr.0).into(),
                            to: (&user_2_addr.0).into(),
                            amount: get_asset("99.0000 GRAEL"),
                            memo: vec![],
                            script: user_1_addr.0.clone().into(),
                        };
                        tx.append_sign(&user_1_addr);
                        TxVariant::TransferTx(tx)
                    };
                    minter
                        .request(MsgRequest::Broadcast(tx))
                        .map(|res| (res, minter))
                }
            })
            .and_then(|(res, minter)| {
                assert_eq!(res, MsgResponse::Broadcast());
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let user_1_bal = minter.chain().get_balance(&user_1_addr.0.into(), &[]);
                assert_eq!(user_1_bal, Some(get_asset("0.0000 GRAEL")));

                let user_2_bal = minter.chain().get_balance(&user_2_addr.0.into(), &[]);
                assert_eq!(user_2_bal, Some(get_asset("99.0000 GRAEL")));

                let minter_addr = ScriptHash::from(&minter.genesis_info().script);
                let minter_bal = minter.chain().get_balance(&minter_addr, &[]);
                // The fee loops back to the minter in a reward transaction generated by the
                // minter block production
                assert_eq!(minter_bal, Some(get_asset("901.0000 GRAEL")));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn insufficient_balance_caused_by_fee() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_addr = ScriptHash::from(&minter.genesis_info().script);
        let to_addr = KeyPair::gen();
        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1001.0000 GRAEL"),
                from: from_addr.clone(),
                to: (&to_addr.0).into(),
                amount: get_asset("0.0000 GRAEL"),
                memo: vec![],
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.and_then(move |res| {
                assert_eq!(
                    res,
                    MsgResponse::Error(net::ErrorKind::TxValidation(
                        verify::TxErr::InsufficientBalance
                    ))
                );
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let chain = minter.chain();
                let cur_bal = chain.get_balance(&to_addr.0.into(), &[]);
                assert_eq!(cur_bal, Some(get_asset("0.0000 GRAEL")));

                let cur_bal = chain.get_balance(&from_addr, &[]);
                assert_eq!(cur_bal, Some(get_asset("1000.0000 GRAEL")));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn insufficient_balance_caused_by_amt() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_addr = ScriptHash::from(&minter.genesis_info().script);
        let to_addr = KeyPair::gen();
        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                from: from_addr.clone(),
                to: (&to_addr.0).into(),
                amount: get_asset("500000.0000 GRAEL"),
                memo: vec![],
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.and_then(move |res| {
                assert_eq!(
                    res,
                    MsgResponse::Error(net::ErrorKind::TxValidation(
                        verify::TxErr::InsufficientBalance
                    ))
                );
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let chain = minter.chain();
                let cur_bal = chain.get_balance(&to_addr.0.into(), &[]);
                assert_eq!(cur_bal, Some(get_asset("0.0000 GRAEL")));

                let cur_bal = chain.get_balance(&from_addr, &[]);
                assert_eq!(cur_bal, Some(get_asset("1000.0000 GRAEL")));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn memo_too_large() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_addr = ScriptHash::from(&minter.genesis_info().script);
        let to_addr = KeyPair::gen();
        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                from: from_addr.clone(),
                to: (&to_addr.0).into(),
                amount: get_asset("1.0000 GRAEL"),
                memo: (0..=godcoin::constants::MAX_MEMO_BYTE_SIZE)
                    .map(|_| 0)
                    .collect(),
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.and_then(move |res| {
                assert_eq!(
                    res,
                    MsgResponse::Error(net::ErrorKind::TxValidation(verify::TxErr::TxTooLarge))
                );
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let chain = minter.chain();
                let cur_bal = chain.get_balance(&to_addr.0.into(), &[]);
                assert_eq!(cur_bal, Some(get_asset("0.0000 GRAEL")));

                let cur_bal = chain.get_balance(&from_addr, &[]);
                assert_eq!(cur_bal, Some(get_asset("1000.0000 GRAEL")));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn script_too_large() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_script = Script::new(
            (0..=godcoin::constants::MAX_SCRIPT_BYTE_SIZE)
                .map(|_| 0)
                .collect(),
        );
        let from_addr = ScriptHash::from(&from_script);
        let to_addr = KeyPair::gen();
        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                from: from_addr,
                to: (&to_addr.0).into(),
                amount: get_asset("1.0000 GRAEL"),
                memo: vec![],
                script: from_script,
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(fut.and_then(move |res| {
            assert_eq!(
                res,
                MsgResponse::Error(net::ErrorKind::TxValidation(verify::TxErr::TxTooLarge))
            );
            System::current().stop();
            Ok(())
        }));
    })
    .unwrap();
}
<commit_msg>Cleanup transfer_from_minter test<commit_after>use actix::prelude::*;
use godcoin::prelude::*;

mod common;
pub use common::*;

#[test]
fn transfer_from_minter() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_addr = ScriptHash::from(&minter.genesis_info().script);
        let from_bal = minter.chain().get_balance(&from_addr, &[]).unwrap();
        let to_addr = KeyPair::gen();
        let amount = get_asset("1.0000 GRAEL");

        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                from: from_addr.clone(),
                to: (&to_addr.0).into(),
                amount,
                memo: vec![],
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.and_then(move |res| {
                assert_eq!(res, MsgResponse::Broadcast());
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let chain = minter.chain();
                let cur_bal = chain.get_balance(&to_addr.0.into(), &[]);
                assert_eq!(cur_bal, Some(amount));

                // The fee transfers back to the minter wallet in the form of a reward tx so it
                // must not be subtracted during the assertion
                let cur_bal = chain.get_balance(&from_addr, &[]);
                assert_eq!(cur_bal, from_bal.sub(amount));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn transfer_from_user() {
    System::run(|| {
        let minter = TestMinter::new();

        let user_1_addr = KeyPair::gen();
        let user_2_addr = KeyPair::gen();

        let fut = {
            let tx = {
                let mut tx = TransferTx {
                    base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                    from: ScriptHash::from(&minter.genesis_info().script),
                    to: (&user_1_addr.0).into(),
                    amount: get_asset("100.0000 GRAEL"),
                    memo: vec![],
                    script: minter.genesis_info().script.clone(),
                };
                tx.append_sign(&minter.genesis_info().wallet_keys[3]);
                tx.append_sign(&minter.genesis_info().wallet_keys[0]);
                TxVariant::TransferTx(tx)
            };
            minter.request(MsgRequest::Broadcast(tx))
        };

        System::current().arbiter().send(
            fut.and_then({
                let user_1_addr = user_1_addr.clone();
                let user_2_addr = user_2_addr.clone();
                move |res| {
                    assert_eq!(res, MsgResponse::Broadcast());
                    let tx = {
                        let mut tx = TransferTx {
                            base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                            from: (&user_1_addr.0).into(),
                            to: (&user_2_addr.0).into(),
                            amount: get_asset("99.0000 GRAEL"),
                            memo: vec![],
                            script: user_1_addr.0.clone().into(),
                        };
                        tx.append_sign(&user_1_addr);
                        TxVariant::TransferTx(tx)
                    };
                    minter
                        .request(MsgRequest::Broadcast(tx))
                        .map(|res| (res, minter))
                }
            })
            .and_then(|(res, minter)| {
                assert_eq!(res, MsgResponse::Broadcast());
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let user_1_bal = minter.chain().get_balance(&user_1_addr.0.into(), &[]);
                assert_eq!(user_1_bal, Some(get_asset("0.0000 GRAEL")));

                let user_2_bal = minter.chain().get_balance(&user_2_addr.0.into(), &[]);
                assert_eq!(user_2_bal, Some(get_asset("99.0000 GRAEL")));

                let minter_addr = ScriptHash::from(&minter.genesis_info().script);
                let minter_bal = minter.chain().get_balance(&minter_addr, &[]);
                // The fee loops back to the minter in a reward transaction generated by the
                // minter block production
                assert_eq!(minter_bal, Some(get_asset("901.0000 GRAEL")));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn insufficient_balance_caused_by_fee() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_addr = ScriptHash::from(&minter.genesis_info().script);
        let to_addr = KeyPair::gen();
        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1001.0000 GRAEL"),
                from: from_addr.clone(),
                to: (&to_addr.0).into(),
                amount: get_asset("0.0000 GRAEL"),
                memo: vec![],
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.and_then(move |res| {
                assert_eq!(
                    res,
                    MsgResponse::Error(net::ErrorKind::TxValidation(
                        verify::TxErr::InsufficientBalance
                    ))
                );
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let chain = minter.chain();
                let cur_bal = chain.get_balance(&to_addr.0.into(), &[]);
                assert_eq!(cur_bal, Some(get_asset("0.0000 GRAEL")));

                let cur_bal = chain.get_balance(&from_addr, &[]);
                assert_eq!(cur_bal, Some(get_asset("1000.0000 GRAEL")));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn insufficient_balance_caused_by_amt() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_addr = ScriptHash::from(&minter.genesis_info().script);
        let to_addr = KeyPair::gen();
        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                from: from_addr.clone(),
                to: (&to_addr.0).into(),
                amount: get_asset("500000.0000 GRAEL"),
                memo: vec![],
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.and_then(move |res| {
                assert_eq!(
                    res,
                    MsgResponse::Error(net::ErrorKind::TxValidation(
                        verify::TxErr::InsufficientBalance
                    ))
                );
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let chain = minter.chain();
                let cur_bal = chain.get_balance(&to_addr.0.into(), &[]);
                assert_eq!(cur_bal, Some(get_asset("0.0000 GRAEL")));

                let cur_bal = chain.get_balance(&from_addr, &[]);
                assert_eq!(cur_bal, Some(get_asset("1000.0000 GRAEL")));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn memo_too_large() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_addr = ScriptHash::from(&minter.genesis_info().script);
        let to_addr = KeyPair::gen();
        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                from: from_addr.clone(),
                to: (&to_addr.0).into(),
                amount: get_asset("1.0000 GRAEL"),
                memo: (0..=godcoin::constants::MAX_MEMO_BYTE_SIZE)
                    .map(|_| 0)
                    .collect(),
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.and_then(move |res| {
                assert_eq!(
                    res,
                    MsgResponse::Error(net::ErrorKind::TxValidation(verify::TxErr::TxTooLarge))
                );
                minter.produce_block().map(|_| minter)
            })
            .and_then(move |minter| {
                let chain = minter.chain();
                let cur_bal = chain.get_balance(&to_addr.0.into(), &[]);
                assert_eq!(cur_bal, Some(get_asset("0.0000 GRAEL")));

                let cur_bal = chain.get_balance(&from_addr, &[]);
                assert_eq!(cur_bal, Some(get_asset("1000.0000 GRAEL")));

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}

#[test]
fn script_too_large() {
    System::run(|| {
        let minter = TestMinter::new();

        let from_script = Script::new(
            (0..=godcoin::constants::MAX_SCRIPT_BYTE_SIZE)
                .map(|_| 0)
                .collect(),
        );
        let from_addr = ScriptHash::from(&from_script);
        let to_addr = KeyPair::gen();
        let tx = {
            let mut tx = TransferTx {
                base: create_tx_header(TxType::TRANSFER, "1.0000 GRAEL"),
                from: from_addr,
                to: (&to_addr.0).into(),
                amount: get_asset("1.0000 GRAEL"),
                memo: vec![],
                script: from_script,
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::TransferTx(tx)
        };
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(fut.and_then(move |res| {
            assert_eq!(
                res,
                MsgResponse::Error(net::ErrorKind::TxValidation(verify::TxErr::TxTooLarge))
            );
            System::current().stop();
            Ok(())
        }));
    })
    .unwrap();
}
<|endoftext|>


<commit_before>use std::net::Ipv4Addr;
use types::*;

/// Bitmasks for the three bit flags field in IPv4
pub mod flags {
    use types::*;

    /// A bitmask with a one in the "Reserved" position.
    pub static RESERVED: u3 = 0b100;
    /// A bitmask with a one in the "Don't fragment" position.
    pub static DF: u3 = 0b010;
    /// A bitmask with a one in the "More fragments" position.
    pub static MF: u3 = 0b001;
}

packet!(Ipv4Packet, MutIpv4Packet, 20);

getters!(Ipv4Packet
    pub fn version(&self) -> u4 {
        read_offset!(self.0, 0, u8) >> 4
    }

    pub fn header_length(&self) -> u4 {
        read_offset!(self.0, 0, u8) & 0x0f
    }

    pub fn dscp(&self) -> u6 {
        read_offset!(self.0, 1, u8) >> 2
    }

    pub fn ecn(&self) -> u2 {
        read_offset!(self.0, 1, u8) & 0x03
    }

    pub fn total_length(&self) -> u16 {
        read_offset!(self.0, 2, u16, from_be)
    }

    pub fn identification(&self) -> u16 {
        read_offset!(self.0, 4, u16, from_be)
    }

    pub fn flags(&self) -> u3 {
        read_offset!(self.0, 6, u8) >> 5
    }

    pub fn dont_fragment(&self) -> bool {
        (self.flags() & flags::DF) != 0
    }

    pub fn more_fragments(&self) -> bool {
        (self.flags() & flags::MF) != 0
    }

    pub fn fragment_offset(&self) -> u13 {
        read_offset!(self.0, 6, u16, from_be) & 0x1fff
    }

    pub fn ttl(&self) -> u8 {
        read_offset!(self.0, 8, u8)
    }

    pub fn protocol(&self) -> Protocol {
        Protocol(read_offset!(self.0, 9, u8))
    }

    pub fn header_checksum(&self) -> u16 {
        read_offset!(self.0, 10, u16, from_be)
    }

    pub fn source(&self) -> Ipv4Addr {
        Ipv4Addr::from(read_offset!(self.0, 12, [u8; 4]))
    }

    pub fn destination(&self) -> Ipv4Addr {
        Ipv4Addr::from(read_offset!(self.0, 16, [u8; 4]))
    }
);

setters!(MutIpv4Packet
    pub fn set_version(&mut self, version: u4) {
        let new_byte = (version << 4) | (read_offset!(self.0, 0, u8) & 0x0f);
        write_offset!(self.0, 0, new_byte, u8);
    }

    pub fn set_header_length(&mut self, header_length: u4) {
        let new_byte = (read_offset!(self.0, 0, u8) & 0xf0) | (header_length & 0x0f);
        write_offset!(self.0, 0, new_byte, u8);
    }

    pub fn set_dscp(&mut self, dscp: u6) {
        let new_byte = (dscp << 2) | (read_offset!(self.0, 1, u8) & 0x03);
        write_offset!(self.0, 1, new_byte, u8);
    }

    pub fn set_ecn(&mut self, ecn: u2) {
        let new_byte = (read_offset!(self.0, 1, u8) & 0xfc) | (ecn & 0x03);
        write_offset!(self.0, 1, new_byte, u8);
    }

    pub fn set_total_length(&mut self, total_length: u16) {
        write_offset!(self.0, 2, total_length, u16, to_be);
    }

    pub fn set_identification(&mut self, identification: u16) {
        write_offset!(self.0, 4, identification, u16, to_be);
    }

    pub fn set_flags(&mut self, flags: u3) {
        let new_byte = (flags << 5) | (read_offset!(self.0, 6, u8) & 0x1f);
        write_offset!(self.0, 6, new_byte, u8);
    }

    pub fn set_fragment_offset(&mut self, fragment_offset: u13) {
        let new_byte = (read_offset!(self.0, 6, u16, from_be) & 0xe000) |
            (fragment_offset & 0x1fff);
        write_offset!(self.0, 6, new_byte, u16, to_be);
    }

    pub fn set_ttl(&mut self, ttl: u8) {
        write_offset!(self.0, 8, ttl, u8);
    }

    pub fn set_protocol(&mut self, protocol: Protocol) {
        write_offset!(self.0, 9, protocol.value(), u8);
    }

    pub fn set_header_checksum(&mut self, checksum: u16) {
        write_offset!(self.0, 10, checksum, u16, to_be);
    }

    pub fn set_source(&mut self, source: Ipv4Addr) {
        write_offset!(self.0, 12, source.octets(), [u8; 4]);
    }

    pub fn set_destination(&mut self, destination: Ipv4Addr) {
        write_offset!(self.0, 16, destination.octets(), [u8; 4]);
    }
);

#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub struct Protocol(pub u8);

impl Protocol {
    #[inline]
    pub fn value(&self) -> u8 {
        self.0
    }
}

pub mod protocols {
    use super::Protocol;

    pub const ICMP: Protocol = Protocol(1);
    pub const TCP: Protocol = Protocol(6);
    pub const UDP: Protocol = Protocol(17);
}

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

    macro_rules! ipv4_setget_test {
        ($name:ident, $set_name:ident, $value:expr, $offset:expr, $expected:expr) => {
            setget_test!(MutIpv4Packet, $name, $set_name, $value, $offset, $expected);
        }
    }

    ipv4_setget_test!(version, set_version, 0xf, 0, [0xf0]);
    ipv4_setget_test!(header_length, set_header_length, 0xf, 0, [0x0f]);
    ipv4_setget_test!(dscp, set_dscp, 0x3f, 1, [0xfc]);
    ipv4_setget_test!(ecn, set_ecn, 0x3, 1, [0x3]);
    ipv4_setget_test!(total_length, set_total_length, 0xffbf, 2, [0xff, 0xbf]);
    ipv4_setget_test!(identification, set_identification, 0xffaf, 4, [0xff, 0xaf]);
    ipv4_setget_test!(flags, set_flags, 0b111, 6, [0xe0]);
    ipv4_setget_test!(
        fragment_offset,
        set_fragment_offset,
        0x1faf,
        6,
        [0x1f, 0xaf]
    );
    ipv4_setget_test!(ttl, set_ttl, 0xff, 8, [0xff]);
    ipv4_setget_test!(protocol, set_protocol, Protocol(0xff), 9, [0xff]);
    ipv4_setget_test!(
        header_checksum,
        set_header_checksum,
        0xfeff,
        10,
        [0xfe, 0xff]
    );
    ipv4_setget_test!(
        source,
        set_source,
        Ipv4Addr::new(192, 168, 15, 1),
        12,
        [192, 168, 15, 1]
    );
    ipv4_setget_test!(
        destination,
        set_destination,
        Ipv4Addr::new(168, 254, 99, 88),
        16,
        [168, 254, 99, 88]
    );

    #[test]
    fn getters_alternating_bits() {
        let backing_data = [0b1010_1010; 20];
        let testee = Ipv4Packet::new(&backing_data).unwrap();
        assert_eq!(0b1010, testee.version());
        assert_eq!(0b1010, testee.header_length());
        assert_eq!(0b101010, testee.dscp());
        assert_eq!(0b10, testee.ecn());
        assert_eq!(0b1010_1010_1010_1010, testee.total_length());
        assert_eq!(0b1010_1010_1010_1010, testee.identification());
        assert_eq!(0b101, testee.flags());
        assert!(!testee.dont_fragment());
        assert!(testee.more_fragments());
        assert_eq!(0b0_1010_1010_1010, testee.fragment_offset());
    }
}
<commit_msg>Move bitfields<commit_after>use std::net::Ipv4Addr;
use types::*;

packet!(Ipv4Packet, MutIpv4Packet, 20);

getters!(Ipv4Packet
    pub fn version(&self) -> u4 {
        read_offset!(self.0, 0, u8) >> 4
    }

    pub fn header_length(&self) -> u4 {
        read_offset!(self.0, 0, u8) & 0x0f
    }

    pub fn dscp(&self) -> u6 {
        read_offset!(self.0, 1, u8) >> 2
    }

    pub fn ecn(&self) -> u2 {
        read_offset!(self.0, 1, u8) & 0x03
    }

    pub fn total_length(&self) -> u16 {
        read_offset!(self.0, 2, u16, from_be)
    }

    pub fn identification(&self) -> u16 {
        read_offset!(self.0, 4, u16, from_be)
    }

    pub fn flags(&self) -> u3 {
        read_offset!(self.0, 6, u8) >> 5
    }

    pub fn dont_fragment(&self) -> bool {
        (self.flags() & flags::DF) != 0
    }

    pub fn more_fragments(&self) -> bool {
        (self.flags() & flags::MF) != 0
    }

    pub fn fragment_offset(&self) -> u13 {
        read_offset!(self.0, 6, u16, from_be) & 0x1fff
    }

    pub fn ttl(&self) -> u8 {
        read_offset!(self.0, 8, u8)
    }

    pub fn protocol(&self) -> Protocol {
        Protocol(read_offset!(self.0, 9, u8))
    }

    pub fn header_checksum(&self) -> u16 {
        read_offset!(self.0, 10, u16, from_be)
    }

    pub fn source(&self) -> Ipv4Addr {
        Ipv4Addr::from(read_offset!(self.0, 12, [u8; 4]))
    }

    pub fn destination(&self) -> Ipv4Addr {
        Ipv4Addr::from(read_offset!(self.0, 16, [u8; 4]))
    }
);

setters!(MutIpv4Packet
    pub fn set_version(&mut self, version: u4) {
        let new_byte = (version << 4) | (read_offset!(self.0, 0, u8) & 0x0f);
        write_offset!(self.0, 0, new_byte, u8);
    }

    pub fn set_header_length(&mut self, header_length: u4) {
        let new_byte = (read_offset!(self.0, 0, u8) & 0xf0) | (header_length & 0x0f);
        write_offset!(self.0, 0, new_byte, u8);
    }

    pub fn set_dscp(&mut self, dscp: u6) {
        let new_byte = (dscp << 2) | (read_offset!(self.0, 1, u8) & 0x03);
        write_offset!(self.0, 1, new_byte, u8);
    }

    pub fn set_ecn(&mut self, ecn: u2) {
        let new_byte = (read_offset!(self.0, 1, u8) & 0xfc) | (ecn & 0x03);
        write_offset!(self.0, 1, new_byte, u8);
    }

    pub fn set_total_length(&mut self, total_length: u16) {
        write_offset!(self.0, 2, total_length, u16, to_be);
    }

    pub fn set_identification(&mut self, identification: u16) {
        write_offset!(self.0, 4, identification, u16, to_be);
    }

    pub fn set_flags(&mut self, flags: u3) {
        let new_byte = (flags << 5) | (read_offset!(self.0, 6, u8) & 0x1f);
        write_offset!(self.0, 6, new_byte, u8);
    }

    pub fn set_fragment_offset(&mut self, fragment_offset: u13) {
        let new_byte = (read_offset!(self.0, 6, u16, from_be) & 0xe000) |
            (fragment_offset & 0x1fff);
        write_offset!(self.0, 6, new_byte, u16, to_be);
    }

    pub fn set_ttl(&mut self, ttl: u8) {
        write_offset!(self.0, 8, ttl, u8);
    }

    pub fn set_protocol(&mut self, protocol: Protocol) {
        write_offset!(self.0, 9, protocol.value(), u8);
    }

    pub fn set_header_checksum(&mut self, checksum: u16) {
        write_offset!(self.0, 10, checksum, u16, to_be);
    }

    pub fn set_source(&mut self, source: Ipv4Addr) {
        write_offset!(self.0, 12, source.octets(), [u8; 4]);
    }

    pub fn set_destination(&mut self, destination: Ipv4Addr) {
        write_offset!(self.0, 16, destination.octets(), [u8; 4]);
    }
);

/// Bitmasks for the three bit flags field in IPv4
pub mod flags {
    use types::*;

    /// A bitmask with a one in the "Reserved" position.
    pub static RESERVED: u3 = 0b100;
    /// A bitmask with a one in the "Don't fragment" position.
    pub static DF: u3 = 0b010;
    /// A bitmask with a one in the "More fragments" position.
    pub static MF: u3 = 0b001;
}

#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub struct Protocol(pub u8);

impl Protocol {
    #[inline]
    pub fn value(&self) -> u8 {
        self.0
    }
}

pub mod protocols {
    use super::Protocol;

    pub const ICMP: Protocol = Protocol(1);
    pub const TCP: Protocol = Protocol(6);
    pub const UDP: Protocol = Protocol(17);
}

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

    macro_rules! ipv4_setget_test {
        ($name:ident, $set_name:ident, $value:expr, $offset:expr, $expected:expr) => {
            setget_test!(MutIpv4Packet, $name, $set_name, $value, $offset, $expected);
        }
    }

    ipv4_setget_test!(version, set_version, 0xf, 0, [0xf0]);
    ipv4_setget_test!(header_length, set_header_length, 0xf, 0, [0x0f]);
    ipv4_setget_test!(dscp, set_dscp, 0x3f, 1, [0xfc]);
    ipv4_setget_test!(ecn, set_ecn, 0x3, 1, [0x3]);
    ipv4_setget_test!(total_length, set_total_length, 0xffbf, 2, [0xff, 0xbf]);
    ipv4_setget_test!(identification, set_identification, 0xffaf, 4, [0xff, 0xaf]);
    ipv4_setget_test!(flags, set_flags, 0b111, 6, [0xe0]);
    ipv4_setget_test!(
        fragment_offset,
        set_fragment_offset,
        0x1faf,
        6,
        [0x1f, 0xaf]
    );
    ipv4_setget_test!(ttl, set_ttl, 0xff, 8, [0xff]);
    ipv4_setget_test!(protocol, set_protocol, Protocol(0xff), 9, [0xff]);
    ipv4_setget_test!(
        header_checksum,
        set_header_checksum,
        0xfeff,
        10,
        [0xfe, 0xff]
    );
    ipv4_setget_test!(
        source,
        set_source,
        Ipv4Addr::new(192, 168, 15, 1),
        12,
        [192, 168, 15, 1]
    );
    ipv4_setget_test!(
        destination,
        set_destination,
        Ipv4Addr::new(168, 254, 99, 88),
        16,
        [168, 254, 99, 88]
    );

    #[test]
    fn getters_alternating_bits() {
        let backing_data = [0b1010_1010; 20];
        let testee = Ipv4Packet::new(&backing_data).unwrap();
        assert_eq!(0b1010, testee.version());
        assert_eq!(0b1010, testee.header_length());
        assert_eq!(0b101010, testee.dscp());
        assert_eq!(0b10, testee.ecn());
        assert_eq!(0b1010_1010_1010_1010, testee.total_length());
        assert_eq!(0b1010_1010_1010_1010, testee.identification());
        assert_eq!(0b101, testee.flags());
        assert!(!testee.dont_fragment());
        assert!(testee.more_fragments());
        assert_eq!(0b0_1010_1010_1010, testee.fragment_offset());
    }
}
<|endoftext|>


<commit_before>use anyhow::Result;
use log::debug;
use thiserror::Error;

use crate::{
    aspace::RelativeAddressSpace,
    module::{Arch, Module, Permissions, Section},
    util, VA,
};

#[derive(Error, Debug)]
pub enum PEError {
    #[error("format not supported: {0}")]
    FormatNotSupported(String),

    #[error("malformed PE file: {0}")]
    MalformedPEFile(String),
}

/// A parsed and loaded PE file.
/// The `pe` field contains the parsed data, courtesy of goblin.
/// The `module` field contains an address space as the PE would be loaded.
///
/// The struct has a reference to the lifetime of the underlying data that's
/// parsed into the PE.
pub struct PE<'a> {
    pub pe:     goblin::pe::PE<'a>,
    pub module: Module,
}

impl<'a> PE<'a> {
    pub fn executable_sections<'b>(self: &'b PE<'a>) -> Box<dyn Iterator<Item = &Section> + 'b> {
        Box::new(
            self.module
                .sections
                .iter()
                .filter(|section| section.perms.intersects(Permissions::X)),
        )
    }
}

fn get_pe(buf: &[u8]) -> Result<goblin::pe::PE> {
    match goblin::Object::parse(buf)? {
        goblin::Object::PE(pe) => Ok(pe),
        goblin::Object::Elf(_) => Err(PEError::FormatNotSupported("elf".to_string()).into()),
        goblin::Object::Archive(_) => Err(PEError::FormatNotSupported("archive".to_string()).into()),
        goblin::Object::Mach(_) => Err(PEError::FormatNotSupported("macho".to_string()).into()),
        goblin::Object::Unknown(_) => Err(PEError::FormatNotSupported("unknown".to_string()).into()),
    }
}

fn load_pe_header(buf: &[u8], pe: &goblin::pe::PE, base_address: VA) -> Result<Section> {
    let hdr_raw_size = match pe.header.optional_header {
        Some(opt) => opt.windows_fields.size_of_headers,
        // assumption: header is at most 0x200 bytes.
        _ => 0x200,
    };

    //   on disk:
    //
    //   +---------------------------------+
    //   |   header        |  sections...  |
    //   +---------------------------------+
    //   .                  \
    //   .  in memory:       \
    //   .                    \
    //   +-----------------+---+        +-------------
    //   |   header        |   |        |  sections...
    //   +-----------------+---+        +-------------
    //                     ^   ^
    //                     |   +--- virtual size
    //                     |        aligned to 0x200
    //                     +-- raw size
    //                         no alignment

    let hdr_raw_size = std::cmp::min(hdr_raw_size as usize, buf.len());
    let hdr_virt_size = util::align(hdr_raw_size as u64, 0x200);

    Ok(Section {
        physical_range: std::ops::Range {
            start: 0x0,
            end:   hdr_raw_size as u64,
        },
        virtual_range:  std::ops::Range {
            start: base_address,
            end:   base_address + hdr_virt_size,
        },
        perms:          Permissions::R,
        name:           "header".to_string(),
    })
}

/// The section can be executed as code.
const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;

/// The section can be read.
const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;

/// The section can be written to.
const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

fn load_pe_section(base_address: VA, section: &goblin::pe::section_table::SectionTable) -> Result<Section> {
    let name = String::from_utf8_lossy(&section.name[..])
        .into_owned()
        .trim_end_matches('\u{0}')
        .trim_end()
        .splitn(2, '\u{0}')
        .next()
        .unwrap()
        .to_string();

    let virtual_size = util::align(section.virtual_size as u64, 0x200);

    let mut perms = Permissions::empty();
    if section.characteristics & IMAGE_SCN_MEM_READ > 0 {
        perms.insert(Permissions::R);
    }
    if section.characteristics & IMAGE_SCN_MEM_WRITE > 0 {
        perms.insert(Permissions::W);
    }
    if section.characteristics & IMAGE_SCN_MEM_EXECUTE > 0 {
        perms.insert(Permissions::X);
    }

    debug!(
        "pe: section: {} at {:#x}",
        name,
        base_address + section.virtual_address as u64
    );

    Ok(Section {
        physical_range: std::ops::Range {
            start: section.pointer_to_raw_data as u64,
            end:   (section.pointer_to_raw_data + section.size_of_raw_data) as u64,
        },
        virtual_range: std::ops::Range {
            start: base_address + section.virtual_address as u64,
            end:   base_address + section.virtual_address as u64 + virtual_size,
        },
        perms,
        name,
    })
}

pub fn load_pe(buf: &[u8]) -> Result<PE> {
    let pe = get_pe(buf)?;

    let arch = match pe.is_64 {
        false => Arch::X32,
        true => Arch::X64,
    };
    debug!("pe: arch: {:?}", arch);

    let base_address = match pe.header.optional_header {
        Some(opt) => opt.windows_fields.image_base,
        _ => {
            debug!("pe: base address: using default: 0x40:000");
            0x40_000
        }
    };
    debug!("pe: base address: {:#x}", base_address);

    let mut sections = vec![load_pe_header(buf, &pe, base_address)?];
    for section in pe.sections.iter() {
        sections.push(load_pe_section(base_address, section)?);
    }

    let max_address = sections.iter().map(|sec| sec.virtual_range.end).max().unwrap();
    let max_page_address = util::align(max_address as u64, 0x1000) - base_address;
    debug!("pe: address space: capacity: {:#x}", max_page_address);

    let mut address_space = RelativeAddressSpace::with_capacity(max_page_address);

    for section in sections.iter() {
        let pstart = section.physical_range.start as usize;
        let pend = section.physical_range.end as usize;
        let pbuf = &buf[pstart..pend];

        // the section range contains VAs,
        // while we're writing to the RelativeAddressSpace.
        // so shift down by `base_address`.
        let vstart = section.virtual_range.start;
        let rstart = vstart - base_address;
        let vsize = util::align(pbuf.len() as u64, 0x1000);
        let vend = vstart + vsize;

        address_space.map.writezx(rstart, pbuf)?;

        debug!(
            "pe: address space: mapped {:#x} - {:#x} {:?}",
            vstart, vend, section.perms
        );
    }

    let module = Module {
        arch,
        sections,
        address_space: address_space.into_absolute(base_address)?,
    };

    debug!("pe: loaded");
    Ok(PE { pe, module })
}

#[cfg(test)]
mod tests {
    use anyhow::Result;

    use crate::{aspace::AddressSpace, rsrc::*};

    #[test]
    fn base_address() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x1_8000_0000, pe.module.address_space.base_address);

        Ok(())
    }

    #[test]
    fn mz_header() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::load_pe(&buf)?;

        // relative read
        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        // absolute read
        assert_eq!(0x4d, pe.module.address_space.read_u8(0x1_8000_0000 + 0x0)?);
        assert_eq!(0x5a, pe.module.address_space.read_u8(0x1_8000_0000 + 0x1)?);

        Ok(())
    }

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn tiny() -> Result<()> {
        let buf = get_buf(Rsrc::TINY);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn nop() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn mimi() -> Result<()> {
        let buf = get_buf(Rsrc::MIMI);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }
}
<commit_msg>pe: don't support .NET assemblies<commit_after>use anyhow::Result;
use log::debug;
use thiserror::Error;

use crate::{
    aspace::RelativeAddressSpace,
    module::{Arch, Module, Permissions, Section},
    util, VA,
};

#[derive(Error, Debug)]
pub enum PEError {
    #[error("format not supported: {0}")]
    FormatNotSupported(String),

    #[error("malformed PE file: {0}")]
    MalformedPEFile(String),
}

/// A parsed and loaded PE file.
/// The `pe` field contains the parsed data, courtesy of goblin.
/// The `module` field contains an address space as the PE would be loaded.
///
/// The struct has a reference to the lifetime of the underlying data that's
/// parsed into the PE.
pub struct PE<'a> {
    pub pe:     goblin::pe::PE<'a>,
    pub module: Module,
}

impl<'a> PE<'a> {
    pub fn executable_sections<'b>(self: &'b PE<'a>) -> Box<dyn Iterator<Item = &Section> + 'b> {
        Box::new(
            self.module
                .sections
                .iter()
                .filter(|section| section.perms.intersects(Permissions::X)),
        )
    }
}

fn get_pe(buf: &[u8]) -> Result<goblin::pe::PE> {
    match goblin::Object::parse(buf)? {
        goblin::Object::PE(pe) => {
            if let Some(opt) = pe.header.optional_header {
                if opt.data_directories.get_clr_runtime_header().is_some() {
                    return Err(PEError::FormatNotSupported(".NET assembly".to_string()).into());
                }
            }
            Ok(pe)
        }
        goblin::Object::Elf(_) => Err(PEError::FormatNotSupported("elf".to_string()).into()),
        goblin::Object::Archive(_) => Err(PEError::FormatNotSupported("archive".to_string()).into()),
        goblin::Object::Mach(_) => Err(PEError::FormatNotSupported("macho".to_string()).into()),
        goblin::Object::Unknown(_) => Err(PEError::FormatNotSupported("unknown".to_string()).into()),
    }
}

fn load_pe_header(buf: &[u8], pe: &goblin::pe::PE, base_address: VA) -> Result<Section> {
    let hdr_raw_size = match pe.header.optional_header {
        Some(opt) => opt.windows_fields.size_of_headers,
        // assumption: header is at most 0x200 bytes.
        _ => 0x200,
    };

    //   on disk:
    //
    //   +---------------------------------+
    //   |   header        |  sections...  |
    //   +---------------------------------+
    //   .                  \
    //   .  in memory:       \
    //   .                    \
    //   +-----------------+---+        +-------------
    //   |   header        |   |        |  sections...
    //   +-----------------+---+        +-------------
    //                     ^   ^
    //                     |   +--- virtual size
    //                     |        aligned to 0x200
    //                     +-- raw size
    //                         no alignment

    let hdr_raw_size = std::cmp::min(hdr_raw_size as usize, buf.len());
    let hdr_virt_size = util::align(hdr_raw_size as u64, 0x200);

    Ok(Section {
        physical_range: std::ops::Range {
            start: 0x0,
            end:   hdr_raw_size as u64,
        },
        virtual_range:  std::ops::Range {
            start: base_address,
            end:   base_address + hdr_virt_size,
        },
        perms:          Permissions::R,
        name:           "header".to_string(),
    })
}

/// The section can be executed as code.
const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;

/// The section can be read.
const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;

/// The section can be written to.
const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

fn load_pe_section(base_address: VA, section: &goblin::pe::section_table::SectionTable) -> Result<Section> {
    let name = String::from_utf8_lossy(&section.name[..])
        .into_owned()
        .trim_end_matches('\u{0}')
        .trim_end()
        .splitn(2, '\u{0}')
        .next()
        .unwrap()
        .to_string();

    let virtual_size = util::align(section.virtual_size as u64, 0x200);

    let mut perms = Permissions::empty();
    if section.characteristics & IMAGE_SCN_MEM_READ > 0 {
        perms.insert(Permissions::R);
    }
    if section.characteristics & IMAGE_SCN_MEM_WRITE > 0 {
        perms.insert(Permissions::W);
    }
    if section.characteristics & IMAGE_SCN_MEM_EXECUTE > 0 {
        perms.insert(Permissions::X);
    }

    debug!(
        "pe: section: {} at {:#x}",
        name,
        base_address + section.virtual_address as u64
    );

    Ok(Section {
        physical_range: std::ops::Range {
            start: section.pointer_to_raw_data as u64,
            end:   (section.pointer_to_raw_data + section.size_of_raw_data) as u64,
        },
        virtual_range: std::ops::Range {
            start: base_address + section.virtual_address as u64,
            end:   base_address + section.virtual_address as u64 + virtual_size,
        },
        perms,
        name,
    })
}

pub fn load_pe(buf: &[u8]) -> Result<PE> {
    let pe = get_pe(buf)?;

    let arch = match pe.is_64 {
        false => Arch::X32,
        true => Arch::X64,
    };
    debug!("pe: arch: {:?}", arch);

    let base_address = match pe.header.optional_header {
        Some(opt) => opt.windows_fields.image_base,
        _ => {
            debug!("pe: base address: using default: 0x40:000");
            0x40_000
        }
    };
    debug!("pe: base address: {:#x}", base_address);

    let mut sections = vec![load_pe_header(buf, &pe, base_address)?];
    for section in pe.sections.iter() {
        sections.push(load_pe_section(base_address, section)?);
    }

    let max_address = sections.iter().map(|sec| sec.virtual_range.end).max().unwrap();
    let max_page_address = util::align(max_address as u64, 0x1000) - base_address;
    debug!("pe: address space: capacity: {:#x}", max_page_address);

    let mut address_space = RelativeAddressSpace::with_capacity(max_page_address);

    for section in sections.iter() {
        let pstart = section.physical_range.start as usize;
        let pend = section.physical_range.end as usize;
        let pbuf = &buf[pstart..pend];

        // the section range contains VAs,
        // while we're writing to the RelativeAddressSpace.
        // so shift down by `base_address`.
        let vstart = section.virtual_range.start;
        let rstart = vstart - base_address;
        let vsize = util::align(pbuf.len() as u64, 0x1000);
        let vend = vstart + vsize;

        address_space.map.writezx(rstart, pbuf)?;

        debug!(
            "pe: address space: mapped {:#x} - {:#x} {:?}",
            vstart, vend, section.perms
        );
    }

    let module = Module {
        arch,
        sections,
        address_space: address_space.into_absolute(base_address)?,
    };

    debug!("pe: loaded");
    Ok(PE { pe, module })
}

#[cfg(test)]
mod tests {
    use anyhow::Result;

    use crate::{aspace::AddressSpace, rsrc::*};

    #[test]
    fn base_address() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x1_8000_0000, pe.module.address_space.base_address);

        Ok(())
    }

    #[test]
    fn mz_header() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::load_pe(&buf)?;

        // relative read
        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        // absolute read
        assert_eq!(0x4d, pe.module.address_space.read_u8(0x1_8000_0000 + 0x0)?);
        assert_eq!(0x5a, pe.module.address_space.read_u8(0x1_8000_0000 + 0x1)?);

        Ok(())
    }

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn tiny() -> Result<()> {
        let buf = get_buf(Rsrc::TINY);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn nop() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn mimi() -> Result<()> {
        let buf = get_buf(Rsrc::MIMI);
        let pe = crate::loader::pe::load_pe(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }
}
<|endoftext|>