<commit_before>use std::fs::File;
use std::io::Read;
use std::env;
use std::path::{Path,PathBuf};
use util::{path_exists, is_dir};
use std::fs::{read_dir};
use toml;

// otry is 'option try'
macro_rules! otry {
    ($e:expr) => (match $e { Some(e) => e, None => return None })
}

// converts errors into None
macro_rules! otry2 {
    ($e:expr) => (match $e { Ok(e) => e, Err(e) => { error!("ERROR!: {:?} {} {}", e, file!(), line!()); return None } })
}

fn find_src_via_lockfile(kratename: &str, cargofile: &Path) -> Option<PathBuf> {
    let mut file = otry2!(File::open(cargofile));
    let mut string = String::new();
    otry2!(file.read_to_string(&mut string));
    let mut parser = toml::Parser::new(&string);
    let lock_table = parser.parse().unwrap();

    debug!("find_src_via_lockfile found lock table {:?}", lock_table);

    let t = match lock_table.get("package") {
        Some(&toml::Value::Array(ref t1)) => t1,
        _ => return None
    };

    for item in t {
        if let &toml::Value::Table(ref t) = item {
            if let Some(&toml::Value::String(ref name)) = t.get("name") {
                if name.replace("-", "_") == kratename {
                    debug!("found matching crate {:?}", t);
                    let version = otry!(getstr(t, "version"));
                    let source = otry!(getstr(t, "source"));

                    if Some("registry") == source.split("+").nth(0) {
                        return get_versioned_cratefile(name, &version);
                    } else if Some("git") == source.split("+").nth(0) {
                        let sha1 = otry!(source.split("#").last());
                        let mut d = otry!(get_cargo_rootdir());
                        d.push("git");
                        d.push("checkouts");
                        d = otry!(find_git_src_dir(d, name, &sha1));
                        d.push("src");
                        d.push("lib.rs");
                        return Some(d);
                    }
                }
            }
        }
    }
    None
}

fn get_cargo_rootdir() -> Option<PathBuf> {
    let mut d = otry!(env::home_dir());

    // try multirust first, since people with multirust installed will often still
    // have an old .cargo directory lying around
    d.push(".multirust");
    d.push("default");
    if let Ok(mut multirustdefault) = File::open(&d) {
        let mut s = String::new();
        otry2!(multirustdefault.read_to_string(&mut s));
        d.pop();
        d.push("toolchains");
        d.push(s.trim());
        d.push("cargo");
        debug!("get_cargo_rootdir root is {:?}",d);
        return Some(d)
    }

    d.pop();
    d.pop();
    d.push(".cargo");
    if path_exists(&d) {
        Some(d)
    } else {
        None
    }
}

fn get_versioned_cratefile(kratename: &str, version: &str) -> Option<PathBuf> {
    let mut d = otry!(get_cargo_rootdir());
    d.push("registry");
    d.push("src");
    d = otry!(find_cratesio_src_dir(d));

    // if version=* then search for the first matching folder
    if version == "*" {
        use std::fs::read_dir;
        let mut start_path = d.clone();
        start_path.push(kratename);
        let start_name = start_path.to_str().unwrap();

        if let Ok(reader) = read_dir(d) {
            if let Some(path) = reader
                                .map(|entry| entry.unwrap().path())
                                .find(|path| path.to_str().unwrap().starts_with(start_name)) {
                d = path.clone();
            } else {
                return None;
            }
        } else { return None; }
    } else {
        d.push(kratename.to_string() + "-" + &version);
    }

    d.push("src");
    debug!("crate path {:?}",d);

    // First, check for package name at root (src/kratename/lib.rs)
    d.push(kratename.to_string());
    if let Err(_) = File::open(&d) {
        // It doesn't exist, so assume src/lib.rs
        d.pop();
    }
    d.push("lib.rs");
    if let Err(_) = File::open(&d) {
        return None;
    }

    Some(d)
 }

fn find_src_via_tomlfile(kratename: &str, cargofile: &Path) -> Option<PathBuf> {
    // only look for 'path' references here.
    // We find the git and crates.io stuff via the lockfile

    let mut file = otry2!(File::open(cargofile));
    let mut string = String::new();
    otry2!(file.read_to_string(&mut string));
    let mut parser = toml::Parser::new(&string);
    let table = otry!(parser.parse());

    // is it this lib?  (e.g. you're searching from tests to find the main library crate)
    if let Some(&toml::Value::Table(ref t)) = table.get("lib") {
        if let Some(&toml::Value::String(ref name)) = t.get("name") {
            if name == kratename {
                debug!("found {} as lib entry in Cargo.toml", kratename);
                if let Some(&toml::Value::String(ref pathstr)) = t.get("path") {
                    let p = Path::new(pathstr);
                    let libpath = otry!(cargofile.parent()).join(p);
                    return Some(libpath);
                }
            }
        }
    }

    // otherwise search the dependencies
    let t = match table.get("dependencies") {
        Some(&toml::Value::Table(ref t)) => t,
        _ => return None
    };

    let mut name = kratename;
    let value = if kratename.contains('_') {
        t.iter().find(|&(k, _)| k.replace("-", "_") == name).map(|(k,v)| {
            name = k;
            v
        })
    } else {
        t.get(kratename)
    };

    match value {
        Some(&toml::Value::Table(ref t)) => {
            // local directory
            let relative_path = otry!(getstr(t, "path"));
            return Some(otry!(cargofile.parent())
                        .join(relative_path)
                        .join("src")
                        .join("lib.rs"));
        },
        Some(&toml::Value::String(ref version)) => {
            // versioned crate
            return get_versioned_cratefile(name, version);
        }
        _ => return None
    }
}

fn find_cratesio_src_dir(d: PathBuf) -> Option<PathBuf> {
    for entry in otry2!(read_dir(d)) {
        let path = otry2!(entry).path();
        if is_dir(path.as_path()) {
            if let Some(ref fname) = path.file_name().and_then(|s| s.to_str()) {
                if fname.starts_with("github.com-") {
                    return Some(path.clone());
                }
            }
        }
    }
    None
}

fn find_git_src_dir(d: PathBuf, name: &str, sha1: &str) -> Option<PathBuf> {
    for entry in otry2!(read_dir(d)) {
        let path = otry2!(entry).path();
        if is_dir(path.as_path()) {
            if let Some(ref fname) = path.file_name().and_then(|s| s.to_str()) {
                if fname.starts_with(name) {
                    let mut d = path.clone();

                    // dirname can be the sha1 or master.
                    d.push(sha1);

                    if !is_dir(d.as_path()) {
                        d.pop();
                        d.push("master");
                    }

                    let retval = d.clone();

                    // check that the checkout matches the commit sha1
                    d.push(".git");
                    d.push("refs");
                    d.push("heads");
                    d.push("master");

                    let mut headref = String::new();
                    otry2!(otry2!(File::open(d)).read_to_string(&mut headref));

                    debug!("git headref is {:?}", headref);

                    if headref.ends_with("\n") {
                        headref.pop();
                    }

                    if sha1 == headref {
                        return Some(retval);
                    }
                }
            }
        }
    }
    None
}

fn getstr(t: &toml::Table, k: &str) -> Option<String> {
    match t.get(k) {
        Some(&toml::Value::String(ref s)) => Some(s.clone()),
        _ => None
    }
}

fn find_cargo_tomlfile(currentfile: &Path) -> Option<PathBuf> {
    let mut f = currentfile.to_path_buf();
    f.push("Cargo.toml");
    if path_exists(f.as_path()) {
        return Some(f);
    } else {
        if f.pop() && f.pop() {
            return find_cargo_tomlfile(&f);
        } else {
            None
        }
    }
}

pub fn get_crate_file(kratename: &str, from_path: &Path) -> Option<PathBuf> {
    if let Some(tomlfile) = find_cargo_tomlfile(from_path) {
        // look in the lockfile first, if there is one
        debug!("get_crate_file tomlfile is {:?}", tomlfile);
        let mut lockfile = tomlfile.clone();
        lockfile.pop();
        lockfile.push("Cargo.lock");
        if path_exists(lockfile.as_path()) {
            if let Some(f) = find_src_via_lockfile(kratename, &lockfile) {
                return Some(f);
            }
        }

        // oh, no luck with the lockfile. Try the tomlfile
        return find_src_via_tomlfile(kratename, &tomlfile)
    }
    None
}
<commit_msg>check CARGO_HOME environment variable when looking for Cargo root dir<commit_after>use std::fs::File;
use std::io::Read;
use std::env;
use std::path::{Path,PathBuf};
use util::{path_exists, is_dir};
use std::fs::{read_dir};
use toml;

// otry is 'option try'
macro_rules! otry {
    ($e:expr) => (match $e { Some(e) => e, None => return None })
}

// converts errors into None
macro_rules! otry2 {
    ($e:expr) => (match $e { Ok(e) => e, Err(e) => { error!("ERROR!: {:?} {} {}", e, file!(), line!()); return None } })
}

fn find_src_via_lockfile(kratename: &str, cargofile: &Path) -> Option<PathBuf> {
    let mut file = otry2!(File::open(cargofile));
    let mut string = String::new();
    otry2!(file.read_to_string(&mut string));
    let mut parser = toml::Parser::new(&string);
    let lock_table = parser.parse().unwrap();

    debug!("find_src_via_lockfile found lock table {:?}", lock_table);

    let t = match lock_table.get("package") {
        Some(&toml::Value::Array(ref t1)) => t1,
        _ => return None
    };

    for item in t {
        if let &toml::Value::Table(ref t) = item {
            if let Some(&toml::Value::String(ref name)) = t.get("name") {
                if name.replace("-", "_") == kratename {
                    debug!("found matching crate {:?}", t);
                    let version = otry!(getstr(t, "version"));
                    let source = otry!(getstr(t, "source"));

                    if Some("registry") == source.split("+").nth(0) {
                        return get_versioned_cratefile(name, &version);
                    } else if Some("git") == source.split("+").nth(0) {
                        let sha1 = otry!(source.split("#").last());
                        let mut d = otry!(get_cargo_rootdir());
                        d.push("git");
                        d.push("checkouts");
                        d = otry!(find_git_src_dir(d, name, &sha1));
                        d.push("src");
                        d.push("lib.rs");
                        return Some(d);
                    }
                }
            }
        }
    }
    None
}

fn get_cargo_rootdir() -> Option<PathBuf> {
    match env::var_os("CARGO_HOME") {
        Some(x) =>
        {
            let d = PathBuf::from(x);
            if path_exists(&d) {
                return Some(d)
            } else {
                return None
            };
        },
        None => ()
    };

    let mut d = otry!(env::home_dir());

    // try multirust first, since people with multirust installed will often still
    // have an old .cargo directory lying around
    d.push(".multirust");
    d.push("default");
    if let Ok(mut multirustdefault) = File::open(&d) {
        let mut s = String::new();
        otry2!(multirustdefault.read_to_string(&mut s));
        d.pop();
        d.push("toolchains");
        d.push(s.trim());
        d.push("cargo");
        debug!("get_cargo_rootdir root is {:?}",d);
        return Some(d)
    }

    d.pop();
    d.pop();
    d.push(".cargo");
    if path_exists(&d) {
        Some(d)
    } else {
        None
    }
}

fn get_versioned_cratefile(kratename: &str, version: &str) -> Option<PathBuf> {
    let mut d = otry!(get_cargo_rootdir());
    d.push("registry");
    d.push("src");
    d = otry!(find_cratesio_src_dir(d));

    // if version=* then search for the first matching folder
    if version == "*" {
        use std::fs::read_dir;
        let mut start_path = d.clone();
        start_path.push(kratename);
        let start_name = start_path.to_str().unwrap();

        if let Ok(reader) = read_dir(d) {
            if let Some(path) = reader
                                .map(|entry| entry.unwrap().path())
                                .find(|path| path.to_str().unwrap().starts_with(start_name)) {
                d = path.clone();
            } else {
                return None;
            }
        } else { return None; }
    } else {
        d.push(kratename.to_string() + "-" + &version);
    }

    d.push("src");
    debug!("crate path {:?}",d);

    // First, check for package name at root (src/kratename/lib.rs)
    d.push(kratename.to_string());
    if let Err(_) = File::open(&d) {
        // It doesn't exist, so assume src/lib.rs
        d.pop();
    }
    d.push("lib.rs");
    if let Err(_) = File::open(&d) {
        return None;
    }

    Some(d)
 }

fn find_src_via_tomlfile(kratename: &str, cargofile: &Path) -> Option<PathBuf> {
    // only look for 'path' references here.
    // We find the git and crates.io stuff via the lockfile

    let mut file = otry2!(File::open(cargofile));
    let mut string = String::new();
    otry2!(file.read_to_string(&mut string));
    let mut parser = toml::Parser::new(&string);
    let table = otry!(parser.parse());

    // is it this lib?  (e.g. you're searching from tests to find the main library crate)
    if let Some(&toml::Value::Table(ref t)) = table.get("lib") {
        if let Some(&toml::Value::String(ref name)) = t.get("name") {
            if name == kratename {
                debug!("found {} as lib entry in Cargo.toml", kratename);
                if let Some(&toml::Value::String(ref pathstr)) = t.get("path") {
                    let p = Path::new(pathstr);
                    let libpath = otry!(cargofile.parent()).join(p);
                    return Some(libpath);
                }
            }
        }
    }

    // otherwise search the dependencies
    let t = match table.get("dependencies") {
        Some(&toml::Value::Table(ref t)) => t,
        _ => return None
    };

    let mut name = kratename;
    let value = if kratename.contains('_') {
        t.iter().find(|&(k, _)| k.replace("-", "_") == name).map(|(k,v)| {
            name = k;
            v
        })
    } else {
        t.get(kratename)
    };

    match value {
        Some(&toml::Value::Table(ref t)) => {
            // local directory
            let relative_path = otry!(getstr(t, "path"));
            return Some(otry!(cargofile.parent())
                        .join(relative_path)
                        .join("src")
                        .join("lib.rs"));
        },
        Some(&toml::Value::String(ref version)) => {
            // versioned crate
            return get_versioned_cratefile(name, version);
        }
        _ => return None
    }
}

fn find_cratesio_src_dir(d: PathBuf) -> Option<PathBuf> {
    for entry in otry2!(read_dir(d)) {
        let path = otry2!(entry).path();
        if is_dir(path.as_path()) {
            if let Some(ref fname) = path.file_name().and_then(|s| s.to_str()) {
                if fname.starts_with("github.com-") {
                    return Some(path.clone());
                }
            }
        }
    }
    None
}

fn find_git_src_dir(d: PathBuf, name: &str, sha1: &str) -> Option<PathBuf> {
    for entry in otry2!(read_dir(d)) {
        let path = otry2!(entry).path();
        if is_dir(path.as_path()) {
            if let Some(ref fname) = path.file_name().and_then(|s| s.to_str()) {
                if fname.starts_with(name) {
                    let mut d = path.clone();

                    // dirname can be the sha1 or master.
                    d.push(sha1);

                    if !is_dir(d.as_path()) {
                        d.pop();
                        d.push("master");
                    }

                    let retval = d.clone();

                    // check that the checkout matches the commit sha1
                    d.push(".git");
                    d.push("refs");
                    d.push("heads");
                    d.push("master");

                    let mut headref = String::new();
                    otry2!(otry2!(File::open(d)).read_to_string(&mut headref));

                    debug!("git headref is {:?}", headref);

                    if headref.ends_with("\n") {
                        headref.pop();
                    }

                    if sha1 == headref {
                        return Some(retval);
                    }
                }
            }
        }
    }
    None
}

fn getstr(t: &toml::Table, k: &str) -> Option<String> {
    match t.get(k) {
        Some(&toml::Value::String(ref s)) => Some(s.clone()),
        _ => None
    }
}

fn find_cargo_tomlfile(currentfile: &Path) -> Option<PathBuf> {
    let mut f = currentfile.to_path_buf();
    f.push("Cargo.toml");
    if path_exists(f.as_path()) {
        return Some(f);
    } else {
        if f.pop() && f.pop() {
            return find_cargo_tomlfile(&f);
        } else {
            None
        }
    }
}

pub fn get_crate_file(kratename: &str, from_path: &Path) -> Option<PathBuf> {
    if let Some(tomlfile) = find_cargo_tomlfile(from_path) {
        // look in the lockfile first, if there is one
        debug!("get_crate_file tomlfile is {:?}", tomlfile);
        let mut lockfile = tomlfile.clone();
        lockfile.pop();
        lockfile.push("Cargo.lock");
        if path_exists(lockfile.as_path()) {
            if let Some(f) = find_src_via_lockfile(kratename, &lockfile) {
                return Some(f);
            }
        }

        // oh, no luck with the lockfile. Try the tomlfile
        return find_src_via_tomlfile(kratename, &tomlfile)
    }
    None
}
<|endoftext|>


<commit_before>//! Atom-class implementation.

use std::ops::{BitXor, Sub, Range, RangeInclusive};
use std::iter::FromIterator;
use std::fmt::{self, Debug};

use regex_syntax;
use itertools::Itertools;

use super::{Atom, ByteOrChar, Element, Repetition, Sequence, Union};
use super::{Distance, Step};
use util::set::{self, Contains};

// ----------------------------------------------------------------
// Class

/// Any member of a class of atoms
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
pub enum ClassMember<T: Atom> {
    /// A single atom
    Atom(T),

    /// A range of atoms
    Range(T, T)
}

impl<T: Atom> ClassMember<T> {
    /// Map the class-member's atoms to a different type
    pub fn map_atoms<U, F>(self, f: F) -> ClassMember<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        match self {
            ClassMember::Atom(a) => ClassMember::Atom(f(a)),
            ClassMember::Range(a, b) => ClassMember::Range(f(a), f(b))
        }
    }

    /// Get the number of individual atoms that this class member represents.
    pub fn len(&self) -> usize
        where T: Distance
    {
        match self {
            &ClassMember::Atom(_) => 1,
            &ClassMember::Range(first, last) => last.distance(&first) + 1,
        }
    }

    /// Fetch an iterator over each discrete value contained by this member.
    #[inline]
     pub fn iter(&self) -> impl Iterator<Item=T>
        where T: Step {
        match self {
            &ClassMember::Atom(a) => AtomIter(a...a),
            &ClassMember::Range(a, b) => AtomIter(a...b),
        }
    }
}

impl<T: Atom> set::Contains<T> for ClassMember<T> {
    /// Check if the member is or contains a particular atom.
    fn contains(&self, x: T) -> bool {
        match self {
            &ClassMember::Atom(a) => x == a,
            &ClassMember::Range(a, b) => (a...b).contains(x),
        }
    }
}

/// Iterator over the atoms specified by a class member.
struct AtomIter<T: Atom>(RangeInclusive<T>);

impl<T: Atom> Iterator for AtomIter<T>
        where T: Step
{
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0 {
            RangeInclusive::NonEmpty{start, end} => {
                let next_start = start.increment();
                if next_start > end {
                    self.0 = RangeInclusive::Empty{at: end};
                } else {
                    self.0 = RangeInclusive::NonEmpty{start: next_start, end: end};
                }
                Some(start) },
            RangeInclusive::Empty{..} => None
        }
    }
}

/*impl<T: Atom> Debug for ClassMember<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ClassMember::Atom(a) => write!(f, "{:?}", a),
            ClassMember::Range(first, last) => write!(f, "{:?}-{:?}", first, last)
        }
    }
}*/

impl<T: Atom> From<T> for ClassMember<T> {
    fn from(a: T) -> Self { ClassMember::Atom(a) }
}

impl From<regex_syntax::ClassRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ClassRange> for ClassMember<char> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<u8> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl<T: Atom> From<Range<T>> for ClassMember<T>
    where T: Sub<usize,Output=T> + Sub<T,Output=usize>
{
    fn from(r: Range<T>) -> Self {
        if r.end - r.start > 1 { ClassMember::Range(r.start, r.end - 1usize) }
        else { ClassMember::Atom(r.start) }
    }
}
// ----------------------------------------------------------------

/// A set of atoms and/or ranges of atoms.
#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
pub struct Class<T: Atom> {
    polarity: Polarity,
    members: Vec<ClassMember<T>>
}

// FIXME: [optimize] All operations on a Class<T>, including observer and
// predicate methods, currently require at least O(N) time.  We could probably
// improve this by changing the storage representation.
impl<T: Atom> Class<T> {
    /// Create a new Class instance with specified polarity and members taken
    /// from the given iterator.
    pub fn new<I,U>(p: Polarity, items: I) -> Self
        where I: IntoIterator<Item=U>,
              ClassMember<T>: From<U> {
        Class{polarity: p, members: Vec::from_iter(items.into_iter().map(|x| x.into()))}
    }

    /// Map the atoms in the class to a different atom type.
    pub fn map_atoms<U, F>(self, f: F) -> Class<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        Class::new(self.polarity, self.members.into_iter().map(|m| m.map_atoms(&f)))
    }

    /// Check whether the class would match a particular atom.
    pub fn matches(&self, x: T) -> bool {
        self.has_member(x).bitxor(self.polarity == Polarity::INVERTED)
    }

    /// Check whether a particular atom is a member of this class.
    /// Unlike `matches`, this method does not take the class's polarity
    /// into account.
    pub fn has_member(&self, x: T) -> bool {
        self.members.iter().any(|&range| range.contains(x))
    }

    /// Get the number of members (atoms) in the class.
    ///
    /// Ranges of atoms count as the number of atoms in each range.
    // FIXME: [bug] handle inverted polarity?
    pub fn len(&self) -> usize {
        self.members.iter().map(|m| m.len()).sum()
    }

    /// Get the class's polarity.
    #[inline]
    pub fn polarity(&self) -> Polarity {
        self.polarity
    }

    /// Check if the class contains zero members.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.members.is_empty()
    }

    /// Fetch an iterator over the members specified in the class.
    ///
    /// Note that if the class has inverted polarity, the atoms in these
    /// members will be those which are **not** matched by the class.
    pub fn iter_members<'a>(&'a self) -> impl Iterator<Item=&'a ClassMember<T>> {
        self.members.iter()
    }
    }
}

impl<T: Atom> Debug for Class<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.polarity == Polarity::NORMAL {
            write!(f, "[{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
        } else {
            write!(f, "[^{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
        }
    }
}

impl<'a, T: 'a + Atom> set::Contains<&'a ClassMember<T>> for Class<T> {
    fn contains(&self, cm: &'a ClassMember<T>) -> bool {
        self.iter_members().any(|m| cm == m)
    }
}
impl<T: Atom> set::Contains<T> for Class<T> {
    fn contains(&self, x: T) -> bool {
        self.matches(x)
    }
}

// FIXME [optimize] I suspect that we need some sort of specialized data
// structure or clever algorithm in order to perform this test in a time better
// than O(N·M).
impl<T: Atom> set::IsSubsetOf<Class<T>> for Class<T> {
    fn is_subset_of(&self, other: &Self) -> bool {
        self.iter_members().flat_map(|m| m.iter()).all(|a| other.contains(a))
    }
}
impl<T: Atom> set::IsSubsetOf<Sequence<T>> for Class<T> {
    /// Any single-atom pattern like Class cannot match the same set of inputs
    /// as a Sequence, because a Sequence is guaranteed to contain at least
    /// two elements.
    #[inline]
    fn is_subset_of(&self, _: &Sequence<T>) -> bool {
        false
    }
}

impl<T: Atom> set::IsSubsetOf<Union<T>> for Class<T> {
    /// A class is a subset of a union if all members of the class are subsets
    /// of the union.
    fn is_subset_of(&self, _: &Union<T>) -> bool {
        unimplemented!()
    }
}

impl<T: Atom> set::IsSubsetOf<Repetition<T>> for Class<T> {
    fn is_subset_of(&self, _: &Repetition<T>) -> bool {
        unimplemented!()
    }
}

impl<T: Atom> set::IsSubsetOf<Element<T>> for Class<T> {
    fn is_subset_of(&self, other: &Element<T>) -> bool {
        match other {
            // Anything we can do, Wildcard can do betteeeeer...
            &Element::Wildcard => true,

            // A class *can* be a subset of an atom -- as long as that atom is
            // the only member of the class.
            &Element::Atom(a) => self.contains(a) && self.polarity == Polarity::NORMAL && self.len() == 1,
            &Element::Class(ref d) => self.is_subset_of(d),
            &Element::Tagged{ref element, ..} => self.is_subset_of(&**element),
            &Element::Sequence(ref seq) => self.is_subset_of(seq),
            &Element::Union(ref union) => self.is_subset_of(union),
            &Element::Repeat(ref rep) => self.is_subset_of(rep),
            &Element::Anchor(_) => false,
        }
    }
}

impl<T: Atom, U> FromIterator<U> for Class<T>
where ClassMember<T>: From<U> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item=U> {
        Class{polarity: Polarity::NORMAL, members: Vec::from_iter(iter.into_iter().map(|x| x.into()))}
    }
}

macro_rules! impl_class_from {
    ($T: ty, $from: ty) => {
        impl<'a> From<$from> for Class<$T> {
            fn from(c: $from) -> Self {
                Self::new(Polarity::NORMAL,
                          c.into_iter().map(|cr| if cr.end != cr.start { ClassMember::Range(cr.start.into(), cr.end.into()) }
                                            else { ClassMember::Atom(cr.start.into()) }))
            }
        }
    };
}

impl_class_from!(char, regex_syntax::CharClass);
impl_class_from!(u8, regex_syntax::ByteClass);
impl_class_from!(ByteOrChar, regex_syntax::CharClass);
impl_class_from!(ByteOrChar, regex_syntax::ByteClass);

// ----------------------------------------------------------------

/** Flag type used to differentiate between normal and inverting matches on
 * pattern elements. */
#[derive(Debug,Copy,Clone,PartialEq, PartialOrd, Eq, Ord)]
pub enum Polarity {
    /// Normal match behavior.
    NORMAL,
    /// Invert the match result.
    INVERTED
}
<commit_msg>pattern/class: add `iter` fn for iterating over all atoms<commit_after>//! Atom-class implementation.

use std::ops::{BitXor, Sub, Range, RangeInclusive};
use std::iter::FromIterator;
use std::fmt::{self, Debug};

use regex_syntax;
use itertools::Itertools;

use super::{Atom, ByteOrChar, Element, Repetition, Sequence, Union};
use super::{Distance, Step};
use util::set::{self, Contains};

// ----------------------------------------------------------------
// Class

/// Any member of a class of atoms
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
pub enum ClassMember<T: Atom> {
    /// A single atom
    Atom(T),

    /// A range of atoms
    Range(T, T)
}

impl<T: Atom> ClassMember<T> {
    /// Map the class-member's atoms to a different type
    pub fn map_atoms<U, F>(self, f: F) -> ClassMember<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        match self {
            ClassMember::Atom(a) => ClassMember::Atom(f(a)),
            ClassMember::Range(a, b) => ClassMember::Range(f(a), f(b))
        }
    }

    /// Get the number of individual atoms that this class member represents.
    pub fn len(&self) -> usize
        where T: Distance
    {
        match self {
            &ClassMember::Atom(_) => 1,
            &ClassMember::Range(first, last) => last.distance(&first) + 1,
        }
    }

    /// Fetch an iterator over each discrete value contained by this member.
    #[inline]
     pub fn iter(&self) -> impl Iterator<Item=T>
        where T: Step {
        match self {
            &ClassMember::Atom(a) => AtomIter(a...a),
            &ClassMember::Range(a, b) => AtomIter(a...b),
        }
    }
}

impl<T: Atom> set::Contains<T> for ClassMember<T> {
    /// Check if the member is or contains a particular atom.
    fn contains(&self, x: T) -> bool {
        match self {
            &ClassMember::Atom(a) => x == a,
            &ClassMember::Range(a, b) => (a...b).contains(x),
        }
    }
}

/// Iterator over the atoms specified by a class member.
struct AtomIter<T: Atom>(RangeInclusive<T>);

impl<T: Atom> Iterator for AtomIter<T>
        where T: Step
{
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        match self.0 {
            RangeInclusive::NonEmpty{start, end} => {
                let next_start = start.increment();
                if next_start > end {
                    self.0 = RangeInclusive::Empty{at: end};
                } else {
                    self.0 = RangeInclusive::NonEmpty{start: next_start, end: end};
                }
                Some(start) },
            RangeInclusive::Empty{..} => None
        }
    }
}

/*impl<T: Atom> Debug for ClassMember<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ClassMember::Atom(a) => write!(f, "{:?}", a),
            ClassMember::Range(first, last) => write!(f, "{:?}-{:?}", first, last)
        }
    }
}*/

impl<T: Atom> From<T> for ClassMember<T> {
    fn from(a: T) -> Self { ClassMember::Atom(a) }
}

impl From<regex_syntax::ClassRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<ByteOrChar> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start.into(), cr.end.into())
    }
}

impl From<regex_syntax::ClassRange> for ClassMember<char> {
    fn from(cr: regex_syntax::ClassRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl From<regex_syntax::ByteRange> for ClassMember<u8> {
    fn from(cr: regex_syntax::ByteRange) -> Self {
        ClassMember::Range(cr.start, cr.end)
    }
}

impl<T: Atom> From<Range<T>> for ClassMember<T>
    where T: Sub<usize,Output=T> + Sub<T,Output=usize>
{
    fn from(r: Range<T>) -> Self {
        if r.end - r.start > 1 { ClassMember::Range(r.start, r.end - 1usize) }
        else { ClassMember::Atom(r.start) }
    }
}
// ----------------------------------------------------------------

/// A set of atoms and/or ranges of atoms.
#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
pub struct Class<T: Atom> {
    polarity: Polarity,
    members: Vec<ClassMember<T>>
}

// FIXME: [optimize] All operations on a Class<T>, including observer and
// predicate methods, currently require at least O(N) time.  We could probably
// improve this by changing the storage representation.
impl<T: Atom> Class<T> {
    /// Create a new Class instance with specified polarity and members taken
    /// from the given iterator.
    pub fn new<I,U>(p: Polarity, items: I) -> Self
        where I: IntoIterator<Item=U>,
              ClassMember<T>: From<U> {
        Class{polarity: p, members: Vec::from_iter(items.into_iter().map(|x| x.into()))}
    }

    /// Map the atoms in the class to a different atom type.
    pub fn map_atoms<U, F>(self, f: F) -> Class<U>
        where F: Fn(T) -> U,
              U: Atom
    {
        Class::new(self.polarity, self.members.into_iter().map(|m| m.map_atoms(&f)))
    }

    /// Check whether the class would match a particular atom.
    pub fn matches(&self, x: T) -> bool {
        self.has_member(x).bitxor(self.polarity == Polarity::INVERTED)
    }

    /// Check whether a particular atom is a member of this class.
    /// Unlike `matches`, this method does not take the class's polarity
    /// into account.
    pub fn has_member(&self, x: T) -> bool {
        self.members.iter().any(|&range| range.contains(x))
    }

    /// Get the number of members (atoms) in the class.
    ///
    /// Ranges of atoms count as the number of atoms in each range.
    // FIXME: [bug] handle inverted polarity?
    pub fn len(&self) -> usize {
        self.members.iter().map(|m| m.len()).sum()
    }

    /// Get the class's polarity.
    #[inline]
    pub fn polarity(&self) -> Polarity {
        self.polarity
    }

    /// Check if the class contains zero members.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.members.is_empty()
    }

    /// Fetch an iterator over the members specified in the class.
    ///
    /// Note that if the class has inverted polarity, the atoms in these
    /// members will be those which are **not** matched by the class.
    pub fn iter_members<'a>(&'a self) -> impl Iterator<Item=&'a ClassMember<T>> {
        self.members.iter()
    }

    /// Fetch an iterator over the atoms specified in the class.
    ///
    /// Note that if the class has inverted polarity, these will be those atoms
    /// which are **not** matched by the class.
    pub fn iter<'a>(&'a self) -> impl Iterator<Item=T> + 'a {
        self.members.iter().flat_map(|m| m.iter())
    }
}

impl<T: Atom> Debug for Class<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.polarity == Polarity::NORMAL {
            write!(f, "[{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
        } else {
            write!(f, "[^{}]", self.members.iter().map(|m| format!("{:?}", m)).join(","))
        }
    }
}

impl<'a, T: 'a + Atom> set::Contains<&'a ClassMember<T>> for Class<T> {
    fn contains(&self, cm: &'a ClassMember<T>) -> bool {
        self.iter_members().any(|m| cm == m)
    }
}
impl<T: Atom> set::Contains<T> for Class<T> {
    fn contains(&self, x: T) -> bool {
        self.matches(x)
    }
}

// FIXME [optimize] I suspect that we need some sort of specialized data
// structure or clever algorithm in order to perform this test in a time better
// than O(N·M).
impl<T: Atom> set::IsSubsetOf<Class<T>> for Class<T> {
    fn is_subset_of(&self, other: &Self) -> bool {
        self.iter_members().flat_map(|m| m.iter()).all(|a| other.contains(a))
    }
}
impl<T: Atom> set::IsSubsetOf<Sequence<T>> for Class<T> {
    /// Any single-atom pattern like Class cannot match the same set of inputs
    /// as a Sequence, because a Sequence is guaranteed to contain at least
    /// two elements.
    #[inline]
    fn is_subset_of(&self, _: &Sequence<T>) -> bool {
        false
    }
}

impl<T: Atom> set::IsSubsetOf<Union<T>> for Class<T> {
    /// A class is a subset of a union if all members of the class are subsets
    /// of the union.
    fn is_subset_of(&self, _: &Union<T>) -> bool {
        unimplemented!()
    }
}

impl<T: Atom> set::IsSubsetOf<Repetition<T>> for Class<T> {
    fn is_subset_of(&self, _: &Repetition<T>) -> bool {
        unimplemented!()
    }
}

impl<T: Atom> set::IsSubsetOf<Element<T>> for Class<T> {
    fn is_subset_of(&self, other: &Element<T>) -> bool {
        match other {
            // Anything we can do, Wildcard can do betteeeeer...
            &Element::Wildcard => true,

            // A class *can* be a subset of an atom -- as long as that atom is
            // the only member of the class.
            &Element::Atom(a) => self.contains(a) && self.polarity == Polarity::NORMAL && self.len() == 1,
            &Element::Class(ref d) => self.is_subset_of(d),
            &Element::Tagged{ref element, ..} => self.is_subset_of(&**element),
            &Element::Sequence(ref seq) => self.is_subset_of(seq),
            &Element::Union(ref union) => self.is_subset_of(union),
            &Element::Repeat(ref rep) => self.is_subset_of(rep),
            &Element::Anchor(_) => false,
        }
    }
}

impl<T: Atom, U> FromIterator<U> for Class<T>
where ClassMember<T>: From<U> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item=U> {
        Class{polarity: Polarity::NORMAL, members: Vec::from_iter(iter.into_iter().map(|x| x.into()))}
    }
}

macro_rules! impl_class_from {
    ($T: ty, $from: ty) => {
        impl<'a> From<$from> for Class<$T> {
            fn from(c: $from) -> Self {
                Self::new(Polarity::NORMAL,
                          c.into_iter().map(|cr| if cr.end != cr.start { ClassMember::Range(cr.start.into(), cr.end.into()) }
                                            else { ClassMember::Atom(cr.start.into()) }))
            }
        }
    };
}

impl_class_from!(char, regex_syntax::CharClass);
impl_class_from!(u8, regex_syntax::ByteClass);
impl_class_from!(ByteOrChar, regex_syntax::CharClass);
impl_class_from!(ByteOrChar, regex_syntax::ByteClass);

// ----------------------------------------------------------------

/** Flag type used to differentiate between normal and inverting matches on
 * pattern elements. */
#[derive(Debug,Copy,Clone,PartialEq, PartialOrd, Eq, Ord)]
pub enum Polarity {
    /// Normal match behavior.
    NORMAL,
    /// Invert the match result.
    INVERTED
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with sessions.
//!
//! Sessions are created per action request and should be used as the only way
//! to communicate with the server. Sessions are responsible for handling action
//! errors (when something goes wrong) or notifying the server that the action
//! finished (by sending appropriate status signal).
//!
//! They also keep track of various statistics (such as number of transferred
//! bytes, action runtime, etc.) and stop the execution if they exceed limits
//! for a particular request.

mod demand;
mod error;
mod response;
mod sink;

use std::convert::TryInto;

use log::{error, info};

use crate::action;
use crate::message;
pub use self::demand::{Demand, Header, Payload};
pub use self::error::{Error, ParseError, MissingFieldError};
use self::response::{Response, Status};
pub use self::sink::{Sink};

/// A specialized `Result` type for sessions.
pub type Result<T> = std::result::Result<T, Error>;

/// Object associating a session with particular action request.
///
/// This is just a convenience type used to avoid threading large numbers of
/// parameters through different function calls.
///
/// Note that the payload can be quite large object, so it should not be stored
/// with the session itself. The payload should be consumed as soon as the
/// request is dispatched to a particular handler and parsed to a concrete type.
pub struct Task<'s, S: Session> {
    /// A session that this task should be run on.
    pub session: &'s mut S,
    /// Serialized data of the action request.
    pub payload: Payload,
}

impl<'s, S: Session> Task<'s, S> {

    /// Executes the task with a particular action handler.
    pub fn execute<R, H>(self, handler: H) -> Result<()>
    where
        R: action::Request,
        H: FnOnce(&mut S, R) -> Result<()>,
    {
        let request = self.payload.parse()?;
        handler(self.session, request)
    }
}

/// Processes given message, handling all errors.
///
/// This function takes a message from the GRR server and interprets it as an
/// action request. It creates a new session object and dispatches the request
/// to an appropriate action handler. Once the action finishes, the function
/// will send a final status message, informing the server about a success
/// or a failure.
///
/// Note that if action execution fails, this function deals with all the errors
/// by sending appropriate information to the server (if possible), logging them
/// and failing hard if a critical error (e.g. communication failure) occurred.
pub fn handle<M>(message: M)
where
    M: TryInto<Demand, Error=ParseError>,
{
    let demand = match message.try_into() {
        Ok(demand) => {
            info!("requested to execute the '{}' action", demand.action);
            demand
        }
        Err(error) => {
            error!("failed to parse the message: {}", error);
            return;
        }
    };

    let mut session = Action::from_demand(&demand);

    let result = action::dispatch(&demand.action, Task {
        session: &mut session,
        payload: demand.payload,
    });

    if let Err(ref error) = result {
        error!("failed to execute the '{}' action: {}", demand.action, error);
    } else {
        info!("finished executing the '{}' action", demand.action);
    }

    let message = match session.status(result).try_into() {
        Ok(message) => message,
        Err(error) => {
            // If we cannot encode the final status message, there is nothing
            // we can do to notify the server, as status is responsible for
            // reporting errors. We can only log the error and carry on.
            error!("failed to encode status message: {}", error);
            return;
        }
    };

    message::send(message);
}

/// Abstraction for various kinds of sessions.
pub trait Session {
    /// Sends a reply to the flow that call the action.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a message to a particular sink.
    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a heartbeat signal to the Fleetspeak process.
    fn heartbeat(&mut self) {
        // TODO: Create a real implementation.
    }
}

/// A session type for unrequested action executions.
///
/// Certain kind of actions are executed not only when a server flow decides to
/// do so, but also upon particular kind of events (e.g. the agent's startup).
/// In such cases, when one needs to trigger action execution manually, ad-hoc
/// sessions should be used.
pub struct Adhoc;

impl Session for Adhoc {

    // TODO: Session trait should be probably split into two traits and then
    // make the actions that do not care about the `reply` method implement the
    // simpler one.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where
        R: action::Response,
    {
        error!("attempted to reply to an ad-hoc session, dropping response");
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// A session type for ordinary action requests.
///
/// This is a normal session type that that is associated with some flow on the
/// server. It keeps track of the responses it sends and collects statistics
/// about network and runtime utilization to kill the action if it is needed.
pub struct Action {
    header: Header,
    next_response_id: u64,
}

impl Action {

    /// Constructs a new session for the given `demand` object.
    pub fn from_demand(demand: &Demand) -> Action {
        // Response identifiers that GRR agents use start at 1. Unfortunately,
        // the server uses this assumption (to determine the number of expected
        // responses when status message is received), so we have to follow this
        // behaviour in RRG as well.
        Action {
            header: demand.header.clone(),
            next_response_id: 1,
        }
    }

    /// Wraps an action response to a session-specific response.
    fn wrap<R>(&self, response: R) -> Response<R>
    where
        R: action::Response
    {
        Response {
            session_id: self.header.session_id.clone(),
            request_id: Some(self.header.request_id),
            response_id: Some(self.next_response_id),
            data: response,
        }
    }

    /// Wraps an action result to a session-specific status response.
    ///
    /// Note that this method consumes the session. The reason for this is that
    /// status response should be obtained as the very last response and no
    /// further replies should be possible after that.
    fn status(self, result: Result<()>) -> Status {
        Status {
            session_id: self.header.session_id,
            request_id: self.header.request_id,
            response_id: self.next_response_id,
            result: result,
        }
    }
}

impl Session for Action {

    fn reply<R: action::Response>(&mut self, response: R) -> Result<()> {
        send(self.wrap(response))?;
        self.next_response_id += 1;

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// Sends a session response to the server.
///
/// Note that this function is not exposed on purpose. Actions should send
/// responses through session objects which introduce a layer of safety. `send`
/// is a low-level utility supposed to be used internally.
fn send<R>(response: Response<R>) -> Result<()>
where
    R: action::Response,
{
    let message = response.try_into()?;
    message::send(message);

    Ok(())
}

#[cfg(test)]
pub mod test {
    use std::any::Any;
    use std::collections::HashMap;

    use super::*;

    /// A session type intended to be used in tests.
    ///
    /// Testing actions with normal session objects can be quite hard, since
    /// they communicate with the outside world (through Fleetspeak). Since we
    /// want to keep the tests minimal and not waste resources on unneeded I/O,
    /// using real sessions is not an option.
    ///
    /// Instead, one can use a `Fake` session. It simply accumulates responses
    /// that the action sends and lets the creator inspect them later.
    pub struct Fake {
        replies: Vec<Box<dyn Any>>,
        responses: HashMap<Sink, Vec<Box<dyn Any>>>,
    }

    impl Fake {

        /// Constructs a new fake session.
        pub fn new() -> Fake {
            Fake {
                replies: Vec::new(),
                responses: std::collections::HashMap::new(),
            }
        }

        /// Yields the number of replies that this session sent so far.
        pub fn reply_count(&self) -> usize {
            self.replies.len()
        }

        /// Retrieves a reply corresponding to the given id.
        ///
        /// The identifier corresponding to the first response is 0, the second
        /// one is 1 and so on.
        ///
        /// This method will panic if a reply with the specified `id` does not
        /// exist or if it exists but has a wrong type.
        pub fn reply<R>(&self, id: usize) -> &R
        where
            R: action::Response + 'static,
        {
            match self.replies::<R>().nth(id) {
                Some(reply) => reply,
                None => panic!("no reply #{}", id),
            }
        }

        pub fn replies<R>(&self) -> impl Iterator<Item = &R>
        where
            R: action::Response + 'static
        {
            self.replies.iter().map(|reply| {
                reply.downcast_ref().expect("unexpected reply type")
            })
        }

        /// Yields the number of responses sent so far to the specified sink.
        pub fn response_count(&self, sink: Sink) -> usize {
            match self.responses.get(&sink) {
                Some(responses) => responses.len(),
                None => 0,
            }
        }

        /// Retrieves a response with the given id sent to a particular sink.
        ///
        /// The identifier corresponding to the first response to the particular
        /// sink is 0, to the second one (to the same sink) is 1 and so on.
        ///
        /// This method will panic if a reply with the specified `id` to the
        /// given `sink` does not exist or if it exists but has wrong type.
        pub fn response<R>(&self, sink: Sink, id: usize) -> &R
        where
            R: action::Response + 'static,
        {
            match self.responses::<R>(sink).nth(id) {
                Some(response) => response,
                None => panic!("no response #{} for sink '{:?}'", id, sink),
            }
        }

        pub fn responses<R>(&self, sink: Sink) -> impl Iterator<Item = &R>
        where
            R: action::Response + 'static,
        {
            let responses = match self.responses.get(&sink) {
                Some(responses) => responses,
                None => panic!("no responses for sink '{:?}'", sink),
            };

            responses.iter().map(move |response| match response.downcast_ref() {
                Some(response) => response,
                None => panic!("unexpected response type in sink '{:?}'", sink),
            })
        }
    }

    impl Session for Fake {

        fn reply<R>(&mut self, response: R) -> Result<()>
        where
            R: action::Response + 'static,
        {
            self.replies.push(Box::new(response));

            Ok(())
        }

        fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
        where
            R: action::Response + 'static,
        {
            let responses = self.responses.entry(sink).or_insert_with(Vec::new);
            responses.push(Box::new(response));

            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_fake_reply_count() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(()).unwrap();
            session.reply(()).unwrap();
            session.reply(()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.reply_count(), 3);
    }

    #[test]
    fn test_fake_response_count() {

        // TODO: Extend this test with more sinks (once we have some more sinks
        // defined).

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, ()).unwrap();
            session.send(Sink::STARTUP, ()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.response_count(Sink::STARTUP), 2);
    }

    #[test]
    fn test_fake_reply_correct_response() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(StringResponse::from("foo")).unwrap();
            session.reply(StringResponse::from("bar")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.reply::<StringResponse>(0).0, "foo");
        assert_eq!(session.reply::<StringResponse>(1).0, "bar");
    }

    #[test]
    #[should_panic(expected = "no reply #0")]
    fn test_fake_reply_incorrect_response_id() {

        fn handle<S: Session>(_: &mut S, _: ()) {
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.reply::<()>(0);
    }

    #[test]
    #[should_panic(expected = "unexpected reply type")]
    fn test_fake_reply_incorrect_response_type() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(StringResponse::from("quux")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.reply::<()>(0);
    }

    #[test]
    fn test_fake_response_correct_response() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, StringResponse::from("foo")).unwrap();
            session.send(Sink::STARTUP, StringResponse::from("bar")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        let response_foo = session.response::<StringResponse>(Sink::STARTUP, 0);
        let response_bar = session.response::<StringResponse>(Sink::STARTUP, 1);
        assert_eq!(response_foo.0, "foo");
        assert_eq!(response_bar.0, "bar");
    }

    #[test]
    #[should_panic(expected = "no responses")]
    fn test_fake_response_empty_sink() {

        fn handle<S: Session>(_: &mut S, _: ()) {
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 0);
    }

    #[test]
    #[should_panic(expected = "no response #42")]
    fn test_fake_response_incorrect_response_id() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, ()).unwrap();
            session.send(Sink::STARTUP, ()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 42);
    }

    #[test]
    #[should_panic(expected = "unexpected response type")]
    fn test_fake_response_incorrect_response_type() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, StringResponse::from("quux")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 0);
    }

    struct StringResponse(String);

    impl<S: Into<String>> From<S> for StringResponse {

        fn from(string: S) -> StringResponse {
            StringResponse(string.into())
        }
    }

    impl action::Response for StringResponse {

        const RDF_NAME: Option<&'static str> = Some("RDFString");

        type Proto = String;

        fn into_proto(self) -> String {
            self.0
        }
    }
}
<commit_msg>Improve panic messages for fake session iterators.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with sessions.
//!
//! Sessions are created per action request and should be used as the only way
//! to communicate with the server. Sessions are responsible for handling action
//! errors (when something goes wrong) or notifying the server that the action
//! finished (by sending appropriate status signal).
//!
//! They also keep track of various statistics (such as number of transferred
//! bytes, action runtime, etc.) and stop the execution if they exceed limits
//! for a particular request.

mod demand;
mod error;
mod response;
mod sink;

use std::convert::TryInto;

use log::{error, info};

use crate::action;
use crate::message;
pub use self::demand::{Demand, Header, Payload};
pub use self::error::{Error, ParseError, MissingFieldError};
use self::response::{Response, Status};
pub use self::sink::{Sink};

/// A specialized `Result` type for sessions.
pub type Result<T> = std::result::Result<T, Error>;

/// Object associating a session with particular action request.
///
/// This is just a convenience type used to avoid threading large numbers of
/// parameters through different function calls.
///
/// Note that the payload can be quite large object, so it should not be stored
/// with the session itself. The payload should be consumed as soon as the
/// request is dispatched to a particular handler and parsed to a concrete type.
pub struct Task<'s, S: Session> {
    /// A session that this task should be run on.
    pub session: &'s mut S,
    /// Serialized data of the action request.
    pub payload: Payload,
}

impl<'s, S: Session> Task<'s, S> {

    /// Executes the task with a particular action handler.
    pub fn execute<R, H>(self, handler: H) -> Result<()>
    where
        R: action::Request,
        H: FnOnce(&mut S, R) -> Result<()>,
    {
        let request = self.payload.parse()?;
        handler(self.session, request)
    }
}

/// Processes given message, handling all errors.
///
/// This function takes a message from the GRR server and interprets it as an
/// action request. It creates a new session object and dispatches the request
/// to an appropriate action handler. Once the action finishes, the function
/// will send a final status message, informing the server about a success
/// or a failure.
///
/// Note that if action execution fails, this function deals with all the errors
/// by sending appropriate information to the server (if possible), logging them
/// and failing hard if a critical error (e.g. communication failure) occurred.
pub fn handle<M>(message: M)
where
    M: TryInto<Demand, Error=ParseError>,
{
    let demand = match message.try_into() {
        Ok(demand) => {
            info!("requested to execute the '{}' action", demand.action);
            demand
        }
        Err(error) => {
            error!("failed to parse the message: {}", error);
            return;
        }
    };

    let mut session = Action::from_demand(&demand);

    let result = action::dispatch(&demand.action, Task {
        session: &mut session,
        payload: demand.payload,
    });

    if let Err(ref error) = result {
        error!("failed to execute the '{}' action: {}", demand.action, error);
    } else {
        info!("finished executing the '{}' action", demand.action);
    }

    let message = match session.status(result).try_into() {
        Ok(message) => message,
        Err(error) => {
            // If we cannot encode the final status message, there is nothing
            // we can do to notify the server, as status is responsible for
            // reporting errors. We can only log the error and carry on.
            error!("failed to encode status message: {}", error);
            return;
        }
    };

    message::send(message);
}

/// Abstraction for various kinds of sessions.
pub trait Session {
    /// Sends a reply to the flow that call the action.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a message to a particular sink.
    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a heartbeat signal to the Fleetspeak process.
    fn heartbeat(&mut self) {
        // TODO: Create a real implementation.
    }
}

/// A session type for unrequested action executions.
///
/// Certain kind of actions are executed not only when a server flow decides to
/// do so, but also upon particular kind of events (e.g. the agent's startup).
/// In such cases, when one needs to trigger action execution manually, ad-hoc
/// sessions should be used.
pub struct Adhoc;

impl Session for Adhoc {

    // TODO: Session trait should be probably split into two traits and then
    // make the actions that do not care about the `reply` method implement the
    // simpler one.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where
        R: action::Response,
    {
        error!("attempted to reply to an ad-hoc session, dropping response");
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// A session type for ordinary action requests.
///
/// This is a normal session type that that is associated with some flow on the
/// server. It keeps track of the responses it sends and collects statistics
/// about network and runtime utilization to kill the action if it is needed.
pub struct Action {
    header: Header,
    next_response_id: u64,
}

impl Action {

    /// Constructs a new session for the given `demand` object.
    pub fn from_demand(demand: &Demand) -> Action {
        // Response identifiers that GRR agents use start at 1. Unfortunately,
        // the server uses this assumption (to determine the number of expected
        // responses when status message is received), so we have to follow this
        // behaviour in RRG as well.
        Action {
            header: demand.header.clone(),
            next_response_id: 1,
        }
    }

    /// Wraps an action response to a session-specific response.
    fn wrap<R>(&self, response: R) -> Response<R>
    where
        R: action::Response
    {
        Response {
            session_id: self.header.session_id.clone(),
            request_id: Some(self.header.request_id),
            response_id: Some(self.next_response_id),
            data: response,
        }
    }

    /// Wraps an action result to a session-specific status response.
    ///
    /// Note that this method consumes the session. The reason for this is that
    /// status response should be obtained as the very last response and no
    /// further replies should be possible after that.
    fn status(self, result: Result<()>) -> Status {
        Status {
            session_id: self.header.session_id,
            request_id: self.header.request_id,
            response_id: self.next_response_id,
            result: result,
        }
    }
}

impl Session for Action {

    fn reply<R: action::Response>(&mut self, response: R) -> Result<()> {
        send(self.wrap(response))?;
        self.next_response_id += 1;

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// Sends a session response to the server.
///
/// Note that this function is not exposed on purpose. Actions should send
/// responses through session objects which introduce a layer of safety. `send`
/// is a low-level utility supposed to be used internally.
fn send<R>(response: Response<R>) -> Result<()>
where
    R: action::Response,
{
    let message = response.try_into()?;
    message::send(message);

    Ok(())
}

#[cfg(test)]
pub mod test {
    use std::any::Any;
    use std::collections::HashMap;

    use super::*;

    /// A session type intended to be used in tests.
    ///
    /// Testing actions with normal session objects can be quite hard, since
    /// they communicate with the outside world (through Fleetspeak). Since we
    /// want to keep the tests minimal and not waste resources on unneeded I/O,
    /// using real sessions is not an option.
    ///
    /// Instead, one can use a `Fake` session. It simply accumulates responses
    /// that the action sends and lets the creator inspect them later.
    pub struct Fake {
        replies: Vec<Box<dyn Any>>,
        responses: HashMap<Sink, Vec<Box<dyn Any>>>,
    }

    impl Fake {

        /// Constructs a new fake session.
        pub fn new() -> Fake {
            Fake {
                replies: Vec::new(),
                responses: std::collections::HashMap::new(),
            }
        }

        /// Yields the number of replies that this session sent so far.
        pub fn reply_count(&self) -> usize {
            self.replies.len()
        }

        /// Retrieves a reply corresponding to the given id.
        ///
        /// The identifier corresponding to the first response is 0, the second
        /// one is 1 and so on.
        ///
        /// This method will panic if a reply with the specified `id` does not
        /// exist or if it exists but has a wrong type.
        pub fn reply<R>(&self, id: usize) -> &R
        where
            R: action::Response + 'static,
        {
            match self.replies::<R>().nth(id) {
                Some(reply) => reply,
                None => panic!("no reply #{}", id),
            }
        }

        pub fn replies<R>(&self) -> impl Iterator<Item = &R>
        where
            R: action::Response + 'static
        {
            self.replies.iter().map(|reply| {
                reply.downcast_ref().expect("unexpected reply type")
            })
        }

        /// Yields the number of responses sent so far to the specified sink.
        pub fn response_count(&self, sink: Sink) -> usize {
            match self.responses.get(&sink) {
                Some(responses) => responses.len(),
                None => 0,
            }
        }

        /// Retrieves a response with the given id sent to a particular sink.
        ///
        /// The identifier corresponding to the first response to the particular
        /// sink is 0, to the second one (to the same sink) is 1 and so on.
        ///
        /// This method will panic if a reply with the specified `id` to the
        /// given `sink` does not exist or if it exists but has wrong type.
        pub fn response<R>(&self, sink: Sink, id: usize) -> &R
        where
            R: action::Response + 'static,
        {
            match self.responses::<R>(sink).nth(id) {
                Some(response) => response,
                None => panic!("no response #{} for sink '{:?}'", id, sink),
            }
        }

        pub fn responses<R>(&self, sink: Sink) -> impl Iterator<Item = &R>
        where
            R: action::Response + 'static,
        {
            // Since the empty iterator (as defined in the standard library) is
            // a specific type, it cannot be returned in one branch but not in
            // another branch.
            //
            // Instead, we use the fact that `Option` is an iterator and then we
            // squash it with `Iterator::flat_map`.
            let responses = self.responses.get(&sink)
                .into_iter()
                .flat_map(std::convert::identity);

            responses.map(move |response| match response.downcast_ref() {
                Some(response) => response,
                None => panic!("unexpected response type in sink '{:?}'", sink),
            })
        }
    }

    impl Session for Fake {

        fn reply<R>(&mut self, response: R) -> Result<()>
        where
            R: action::Response + 'static,
        {
            self.replies.push(Box::new(response));

            Ok(())
        }

        fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
        where
            R: action::Response + 'static,
        {
            let responses = self.responses.entry(sink).or_insert_with(Vec::new);
            responses.push(Box::new(response));

            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_fake_reply_count() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(()).unwrap();
            session.reply(()).unwrap();
            session.reply(()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.reply_count(), 3);
    }

    #[test]
    fn test_fake_response_count() {

        // TODO: Extend this test with more sinks (once we have some more sinks
        // defined).

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, ()).unwrap();
            session.send(Sink::STARTUP, ()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.response_count(Sink::STARTUP), 2);
    }

    #[test]
    fn test_fake_reply_correct_response() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(StringResponse::from("foo")).unwrap();
            session.reply(StringResponse::from("bar")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.reply::<StringResponse>(0).0, "foo");
        assert_eq!(session.reply::<StringResponse>(1).0, "bar");
    }

    #[test]
    #[should_panic(expected = "no reply #0")]
    fn test_fake_reply_incorrect_response_id() {

        fn handle<S: Session>(_: &mut S, _: ()) {
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.reply::<()>(0);
    }

    #[test]
    #[should_panic(expected = "unexpected reply type")]
    fn test_fake_reply_incorrect_response_type() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(StringResponse::from("quux")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.reply::<()>(0);
    }

    #[test]
    fn test_fake_response_correct_response() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, StringResponse::from("foo")).unwrap();
            session.send(Sink::STARTUP, StringResponse::from("bar")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        let response_foo = session.response::<StringResponse>(Sink::STARTUP, 0);
        let response_bar = session.response::<StringResponse>(Sink::STARTUP, 1);
        assert_eq!(response_foo.0, "foo");
        assert_eq!(response_bar.0, "bar");
    }

    #[test]
    #[should_panic(expected = "no response #42")]
    fn test_fake_response_incorrect_response_id() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, ()).unwrap();
            session.send(Sink::STARTUP, ()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 42);
    }

    #[test]
    #[should_panic(expected = "unexpected response type")]
    fn test_fake_response_incorrect_response_type() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, StringResponse::from("quux")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 0);
    }

    struct StringResponse(String);

    impl<S: Into<String>> From<S> for StringResponse {

        fn from(string: S) -> StringResponse {
            StringResponse(string.into())
        }
    }

    impl action::Response for StringResponse {

        const RDF_NAME: Option<&'static str> = Some("RDFString");

        type Proto = String;

        fn into_proto(self) -> String {
            self.0
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 Dmytro Milinevskyi <dmilinevskyi@gmail.com>

// 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

// http://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.

extern crate chrono;
use self::chrono::prelude::*;

extern crate time;

extern crate parking_lot;
use self::parking_lot::Mutex;

use std::sync::Arc;

use std::fmt;
use std::fmt::Write;

use formatters::Formatter;
use levels::LogLevel;
use record::Record;

#[doc(hidden)]
#[derive(Clone)]
pub struct RecordMeta {
    pub level: LogLevel,
    pub module: &'static str,
    pub file: &'static str,
    pub line: u32,
}

struct RecordLazyMetaInner {
    msg: Option<Arc<String>>,
    formatted: Option<Arc<String>>,
    ts_utc: Option<Arc<DateTime<UTC>>>,
}

impl RecordLazyMetaInner {
    #[inline(always)]
    fn new() -> Self {
        RecordLazyMetaInner {
            msg: None,
            formatted: None,
            ts_utc: None,
        }
    }

    fn mk_msg<'a>(&mut self, args: fmt::Arguments<'a>) {
        if self.msg.is_none() {
            let mut mstr = String::new();
            mstr.write_fmt(args).unwrap();
            self.msg = Some(Arc::new(mstr));
        }
    }

    fn mk_ts_utc(&mut self, ts: &time::Timespec) {
        if self.ts_utc.is_none() {
            let naive = chrono::NaiveDateTime::from_timestamp(ts.sec, ts.nsec as u32);
            self.ts_utc = Some(Arc::new(chrono::DateTime::from_utc(naive, chrono::UTC)));
        }
    }
}

struct RecordLazyMeta {
    irecord: Mutex<RecordLazyMetaInner>,
    formatter: Arc<Formatter>,
}

impl RecordLazyMeta {
    #[inline(always)]
    fn new(formatter: Arc<Formatter>) -> Self {
        RecordLazyMeta {
            irecord: Mutex::new(RecordLazyMetaInner::new()),
            formatter: formatter,
        }
    }

    fn msg<'a>(&self, args: fmt::Arguments<'a>) -> Arc<String> {
        let mut irecord = self.irecord.lock();
        irecord.mk_msg(args);
        let msg = irecord.msg.as_ref().unwrap();
        msg.clone()
    }

    fn formatted(&self, record: &Record) -> Arc<String> {
        {
            let irecord = self.irecord.lock();
            let format = irecord.formatted.is_none();
            drop(irecord);

            if format {
                let formatted = Arc::new((self.formatter)(record));
                let mut irecord = self.irecord.lock();
                irecord.formatted = Some(formatted.clone());
                return formatted;
            }
        }

        let irecord = self.irecord.lock();
        let formatted = irecord.formatted.as_ref().unwrap();
        formatted.clone()
    }

    fn ts_utc(&self, ts: &time::Timespec) -> Arc<DateTime<UTC>> {
        let mut irecord = self.irecord.lock();
        irecord.mk_ts_utc(ts);
        let ts_utc = irecord.ts_utc.as_ref().unwrap();
        ts_utc.clone()
    }
}

// TODO: use pub(crate) when stabilized (should in v1.18)
// https://github.com/rust-lang/rust/issues/32409
#[doc(hidden)]
pub struct SyncRecord<'a> {
    irecord: &'static RecordMeta,
    args: fmt::Arguments<'a>,
    precord: Arc<RecordLazyMeta>,
    ts: time::Timespec,
}

impl<'a> SyncRecord<'a> {
    // TODO: use pub(crate) when stabilized (should in v1.18)
    // https://github.com/rust-lang/rust/issues/32409
    #[doc(hidden)]
    #[inline(always)]
    pub fn new(record: &'static RecordMeta,
               ts: time::Timespec,
               args: fmt::Arguments<'a>,
               formatter: Arc<Formatter>) -> Self {
        SyncRecord {
            irecord: record,
            ts: ts,
            args: args,
            precord: Arc::new(RecordLazyMeta::new(formatter)),
        }
    }
}

impl<'a> Record for SyncRecord<'a> {
    #[inline(always)]
    fn level(&self) -> LogLevel {
        self.irecord.level
    }

    #[inline(always)]
    fn module(&self) -> &'static str {
        self.irecord.module
    }

    #[inline(always)]
    fn file(&self) -> &'static str {
        self.irecord.file
    }

    #[inline(always)]
    fn line(&self) -> u32 {
        self.irecord.line
    }

    #[inline(always)]
    fn ts(&self) -> time::Timespec {
        self.ts
    }

    fn msg(&self) -> Arc<String> {
        return self.precord.msg(self.args)
    }

    fn formatted(&self) -> Arc<String> {
        return self.precord.formatted(self)
    }

    fn ts_utc(&self) -> Arc<DateTime<UTC>> {
        return self.precord.ts_utc(&self.ts)
    }
}

// TODO: use pub(crate) when stabilized (should in v1.18)
// https://github.com/rust-lang/rust/issues/32409
#[doc(hidden)]
pub struct AsyncRecord {
    irecord: &'static RecordMeta,
    msg: Arc<String>,
    precord: Arc<RecordLazyMeta>,
    ts: time::Timespec,
}

impl Record for AsyncRecord {
    #[inline(always)]
    fn level(&self) -> LogLevel {
        self.irecord.level
    }

    #[inline(always)]
    fn module(&self) -> &'static str {
        self.irecord.module
    }

    #[inline(always)]
    fn file(&self) -> &'static str {
        self.irecord.file
    }

    #[inline(always)]
    fn line(&self) -> u32 {
        self.irecord.line
    }

    #[inline(always)]
    fn ts(&self) -> time::Timespec {
        self.ts
    }

    fn msg(&self) -> Arc<String> {
        self.msg.clone()
    }

    fn formatted(&self) -> Arc<String> {
        return self.precord.formatted(self)
    }

    fn ts_utc(&self) -> Arc<DateTime<UTC>> {
        return self.precord.ts_utc(&self.ts)
    }
}

impl<'a> From<SyncRecord<'a>> for AsyncRecord {
    #[inline(always)]
    fn from(orig: SyncRecord) -> AsyncRecord {
        let mut mstr = String::new();
        mstr.write_fmt(orig.args).unwrap();
        AsyncRecord {
            irecord: orig.irecord,
            msg: Arc::new(mstr),
            precord: orig.precord,
            ts: orig.ts,
        }
    }
}
<commit_msg>record: preallocate some space for the fmt string<commit_after>// Copyright 2017 Dmytro Milinevskyi <dmilinevskyi@gmail.com>

// 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

// http://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.

extern crate chrono;
use self::chrono::prelude::*;

extern crate time;

extern crate parking_lot;
use self::parking_lot::Mutex;

use std::sync::Arc;

use std::fmt;
use std::fmt::Write;

use formatters::Formatter;
use levels::LogLevel;
use record::Record;

const PREALLOC: usize = 128;

#[doc(hidden)]
#[derive(Clone)]
pub struct RecordMeta {
    pub level: LogLevel,
    pub module: &'static str,
    pub file: &'static str,
    pub line: u32,
}

struct RecordLazyMetaInner {
    msg: Option<Arc<String>>,
    formatted: Option<Arc<String>>,
    ts_utc: Option<Arc<DateTime<UTC>>>,
}

impl RecordLazyMetaInner {
    #[inline(always)]
    fn new() -> Self {
        RecordLazyMetaInner {
            msg: None,
            formatted: None,
            ts_utc: None,
        }
    }

    fn mk_msg<'a>(&mut self, args: fmt::Arguments<'a>) {
        if self.msg.is_none() {
            let mut mstr = String::with_capacity(PREALLOC);
            mstr.write_fmt(args).unwrap();
            self.msg = Some(Arc::new(mstr));
        }
    }

    fn mk_ts_utc(&mut self, ts: &time::Timespec) {
        if self.ts_utc.is_none() {
            let naive = chrono::NaiveDateTime::from_timestamp(ts.sec, ts.nsec as u32);
            self.ts_utc = Some(Arc::new(chrono::DateTime::from_utc(naive, chrono::UTC)));
        }
    }
}

struct RecordLazyMeta {
    irecord: Mutex<RecordLazyMetaInner>,
    formatter: Arc<Formatter>,
}

impl RecordLazyMeta {
    #[inline(always)]
    fn new(formatter: Arc<Formatter>) -> Self {
        RecordLazyMeta {
            irecord: Mutex::new(RecordLazyMetaInner::new()),
            formatter: formatter,
        }
    }

    fn msg<'a>(&self, args: fmt::Arguments<'a>) -> Arc<String> {
        let mut irecord = self.irecord.lock();
        irecord.mk_msg(args);
        let msg = irecord.msg.as_ref().unwrap();
        msg.clone()
    }

    fn formatted(&self, record: &Record) -> Arc<String> {
        {
            let irecord = self.irecord.lock();
            let format = irecord.formatted.is_none();
            drop(irecord);

            if format {
                let formatted = Arc::new((self.formatter)(record));
                let mut irecord = self.irecord.lock();
                irecord.formatted = Some(formatted.clone());
                return formatted;
            }
        }

        let irecord = self.irecord.lock();
        let formatted = irecord.formatted.as_ref().unwrap();
        formatted.clone()
    }

    fn ts_utc(&self, ts: &time::Timespec) -> Arc<DateTime<UTC>> {
        let mut irecord = self.irecord.lock();
        irecord.mk_ts_utc(ts);
        let ts_utc = irecord.ts_utc.as_ref().unwrap();
        ts_utc.clone()
    }
}

// TODO: use pub(crate) when stabilized (should in v1.18)
// https://github.com/rust-lang/rust/issues/32409
#[doc(hidden)]
pub struct SyncRecord<'a> {
    irecord: &'static RecordMeta,
    args: fmt::Arguments<'a>,
    precord: Arc<RecordLazyMeta>,
    ts: time::Timespec,
}

impl<'a> SyncRecord<'a> {
    // TODO: use pub(crate) when stabilized (should in v1.18)
    // https://github.com/rust-lang/rust/issues/32409
    #[doc(hidden)]
    #[inline(always)]
    pub fn new(record: &'static RecordMeta,
               ts: time::Timespec,
               args: fmt::Arguments<'a>,
               formatter: Arc<Formatter>) -> Self {
        SyncRecord {
            irecord: record,
            ts: ts,
            args: args,
            precord: Arc::new(RecordLazyMeta::new(formatter)),
        }
    }
}

impl<'a> Record for SyncRecord<'a> {
    #[inline(always)]
    fn level(&self) -> LogLevel {
        self.irecord.level
    }

    #[inline(always)]
    fn module(&self) -> &'static str {
        self.irecord.module
    }

    #[inline(always)]
    fn file(&self) -> &'static str {
        self.irecord.file
    }

    #[inline(always)]
    fn line(&self) -> u32 {
        self.irecord.line
    }

    #[inline(always)]
    fn ts(&self) -> time::Timespec {
        self.ts
    }

    fn msg(&self) -> Arc<String> {
        return self.precord.msg(self.args)
    }

    fn formatted(&self) -> Arc<String> {
        return self.precord.formatted(self)
    }

    fn ts_utc(&self) -> Arc<DateTime<UTC>> {
        return self.precord.ts_utc(&self.ts)
    }
}

// TODO: use pub(crate) when stabilized (should in v1.18)
// https://github.com/rust-lang/rust/issues/32409
#[doc(hidden)]
pub struct AsyncRecord {
    irecord: &'static RecordMeta,
    msg: Arc<String>,
    precord: Arc<RecordLazyMeta>,
    ts: time::Timespec,
}

impl Record for AsyncRecord {
    #[inline(always)]
    fn level(&self) -> LogLevel {
        self.irecord.level
    }

    #[inline(always)]
    fn module(&self) -> &'static str {
        self.irecord.module
    }

    #[inline(always)]
    fn file(&self) -> &'static str {
        self.irecord.file
    }

    #[inline(always)]
    fn line(&self) -> u32 {
        self.irecord.line
    }

    #[inline(always)]
    fn ts(&self) -> time::Timespec {
        self.ts
    }

    fn msg(&self) -> Arc<String> {
        self.msg.clone()
    }

    fn formatted(&self) -> Arc<String> {
        return self.precord.formatted(self)
    }

    fn ts_utc(&self) -> Arc<DateTime<UTC>> {
        return self.precord.ts_utc(&self.ts)
    }
}

impl<'a> From<SyncRecord<'a>> for AsyncRecord {
    #[inline(always)]
    fn from(orig: SyncRecord) -> AsyncRecord {
        let mut mstr = String::with_capacity(PREALLOC);
        mstr.write_fmt(orig.args).unwrap();
        AsyncRecord {
            irecord: orig.irecord,
            msg: Arc::new(mstr),
            precord: orig.precord,
            ts: orig.ts,
        }
    }
}
<|endoftext|>


<commit_before>mod crc;
#[macro_use]
mod macros;
mod types;

pub use self::crc::{crc8, crc16};
pub use self::types::{ErrorKind, ByteStream, ReadStream};

use nom::{self, IResult};
use metadata::{Metadata, metadata_parser};

use std::ops::{Add, AddAssign, BitAnd, BitOr, Mul, Sub, Shl, ShlAssign, Shr};
use std::io;

/// An interface for parsing through some type of producer to a byte stream.
///
/// External parsers get passed in and consumes the bytes held internally
/// and outputs the `Result` of that parser.
pub trait StreamProducer {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T, ErrorKind>;
}

/// An abstraction trait for keeping different sized integers.
pub trait Sample: PartialEq + Eq + Sized + Clone + Copy +
                  Add<Output = Self> + AddAssign +
                  BitAnd<Self, Output = Self> + BitOr<Self, Output = Self> +
                  Mul<Output = Self> + Shl<u32, Output = Self>  +
                  ShlAssign<u32> + Shr<u32, Output = Self> +
                  Shr<i8, Output = Self> + Shr<i32, Output = Self> +
                  Sub<Output = Self> {
  /// The normal size for the current a `Sample`.
  type Normal;

  /// The size, in bits, for the `Sample::Normal`.
  fn size() -> usize;

  /// The size, in bits, for the `Sample`.
  fn size_extended() -> usize;

  /// Convert the extended `Sample` to the normal.
  fn to_normal(sample: Self) -> Option<Self::Normal>;

  /// Convert an i8 into a `Sample`.
  fn from_i8(sample: i8) -> Self;

  /// Convert an i16 into a `Sample`.
  fn from_i16(sample: i16) -> Self;

  /// Convert an i32 into a `Sample`.
  ///
  /// With `Sample` sometimes being smaller than a i32, there is a chance
  /// for this function to return an incorrect number. So when the number is
  /// larger of smaller than the current `Sample`, it returns `None`
  /// otherwise `Some(sample)`.
  fn from_i32(sample: i32) -> Option<Self>;

  /// Convert an i32 into a `Sample`.
  fn from_i32_lossy(sample: i32) -> Self;
}

/// A trait for defining the size of a sample.
pub trait SampleSize {
  /// The internal integer size used with `Stream::iter`.
  ///
  /// Rather than making the user of `Stream::iter` will remember what the
  /// extended size of the output sample will be, this is to map to the
  /// value based on the current integer size used.
  type Extended: Sample;
}

impl SampleSize for i8 {
  type Extended = i16;
}

impl SampleSize for i16 {
  type Extended = i32;
}

impl SampleSize for i32 {
  type Extended = i64;
}

pub trait WriteExtension: io::Write {
  fn write_u8(&mut self, number: u8) -> io::Result<()>;

  fn write_be_u16(&mut self, number: u16) -> io::Result<()>;
  fn write_le_u16(&mut self, number: u16) -> io::Result<()>;

  fn write_be_u24(&mut self, number: u32) -> io::Result<()>;
  fn write_le_u24(&mut self, number: u32) -> io::Result<()>;

  fn write_be_u32(&mut self, number: u32) -> io::Result<()>;
  fn write_le_u32(&mut self, number: u32) -> io::Result<()>;

  fn write_be_u64(&mut self, number: u64) -> io::Result<()>;
  fn write_le_u64(&mut self, number: u64) -> io::Result<()>;
}

impl<Write> WriteExtension for Write where Write: io::Write {
  fn write_u8(&mut self, number: u8) -> io::Result<()> {
    self.write_all(&[number])
  }

  fn write_be_u16(&mut self, number: u16) -> io::Result<()> {
    let mut buffer = [0; 2];

    buffer[0] = (number >> 8) as u8;
    buffer[1] = number as u8;

    self.write_all(&buffer)
  }

  fn write_le_u16(&mut self, number: u16) -> io::Result<()> {
    let mut buffer = [0; 2];

    buffer[2] = number as u8;
    buffer[1] = (number >> 8) as u8;

    self.write_all(&buffer)
  }
}

// Convert one to four byte slices into an unsigned 32-bit number.
//
// NOTE: This assumes big-endian since most numbers in the FLAC binary are
// that endianness.
#[inline]
pub fn to_u32(bytes: &[u8]) -> u32 {
  let length = bytes.len();

  debug_assert!(length <= 4);

  (0..length).fold(0, |result, i|
    result + ((bytes[i] as u32) << ((length - 1 - i) * 8))
  )
}

// Extends a signed value of a specific bit size to a larger bit size.
//
// NOTE: This assumes that the larger bit size will be 32 bit since that is
// the largest sample size supported in FLAC.
pub fn extend_sign(value: u32, bit_count: usize) -> i32 {
  if bit_count >= 32 || value < (1 << (bit_count - 1)) {
    value as i32
  } else {
    (value as i32).wrapping_sub(1 << bit_count)
  }
}

// Bit shifted version for two to the power of a given exponent.
#[inline]
pub fn power_of_two(exponent: u32) -> u32 {
  debug_assert!(exponent <= 31);

  1 << exponent
}

#[derive(PartialEq, Eq)]
enum ParserState {
  Header,
  StreamInfo,
  Metadata
}

fn parser<'a>(input: &'a [u8], state: &mut ParserState)
              -> IResult<&'a [u8], Metadata, ErrorKind> {
  let mut slice = input;
  let error     = nom::Err::Code(nom::ErrorKind::Custom(ErrorKind::Unknown));

  if *state == ParserState::Header {
    let (i, _) = try_parser! {
      to_custom_error!(slice, tag!("fLaC"), HeaderParser)
    };

    slice  = i;
    *state = ParserState::StreamInfo;
  }

  match *state {
    ParserState::StreamInfo => {
      let (i, block) = try_parse!(slice, metadata_parser);

      if block.is_stream_info() {
        *state = ParserState::Metadata;

        IResult::Done(i, block)
      } else {
        IResult::Error(error)
      }
    }
    ParserState::Metadata   => metadata_parser(slice),
    _                       => IResult::Error(error),
  }
}

pub fn many_metadata<S, F>(stream: &mut S, mut f: F) -> Result<(), ErrorKind>
 where S: StreamProducer,
       F: FnMut(Metadata) {
  let mut state  = ParserState::Header;
  let mut result = Ok(());

  loop {
    match stream.parse(|i| parser(i, &mut state)) {
      Ok(block)                => {
        let is_last = block.is_last();

        f(block);

        if is_last {
          break;
        }
      }
      Err(ErrorKind::Continue) => continue,
      Err(e)                   => {
        result = Err(e);

        break;
      }
    }
  }

  result
}

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

  #[test]
  #[should_panic]
  fn test_panic_to_u32() {
    to_u32(&[0x00, 0x00, 0x00, 0x00, 0x00]);
  }

  #[test]
  fn test_to_u32() {
    let bytes = [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef];

    assert_eq!(to_u32(&bytes[0..1]), 0x00000001);
    assert_eq!(to_u32(&bytes[3..5]), 0x00006789);
    assert_eq!(to_u32(&bytes[1..4]), 0x00234567);
    assert_eq!(to_u32(&bytes[4..]), 0x89abcdef);
  }

  #[test]
  fn test_extend_sign() {
    assert_eq!(extend_sign(32, 6), -32);
    assert_eq!(extend_sign(31, 6), 31);
    assert_eq!(extend_sign(128, 8), -128);
    assert_eq!(extend_sign(127, 8), 127);

    assert_eq!(extend_sign(2048, 12), -2048);
    assert_eq!(extend_sign(2047, 12), 2047);
    assert_eq!(extend_sign(32768, 16), -32768);
    assert_eq!(extend_sign(32767, 16), 32767);

    assert_eq!(extend_sign(8388608, 24), -8388608);
    assert_eq!(extend_sign(8388607, 24), 8388607);
    assert_eq!(extend_sign(2147483648, 32), -2147483648);
    assert_eq!(extend_sign(2147483647, 32), 2147483647);
  }

  #[test]
  #[should_panic]
  fn test_panic_power_of_two() {
    power_of_two(32);
  }

  #[test]
  fn test_power_of_two() {
    assert_eq!(power_of_two(0), 1);
    assert_eq!(power_of_two(1), 2);
    assert_eq!(power_of_two(2), 4);
    assert_eq!(power_of_two(10), 1024);
    assert_eq!(power_of_two(31), 2147483648);
  }
}
<commit_msg>Add implementation for writing u24<commit_after>mod crc;
#[macro_use]
mod macros;
mod types;

pub use self::crc::{crc8, crc16};
pub use self::types::{ErrorKind, ByteStream, ReadStream};

use nom::{self, IResult};
use metadata::{Metadata, metadata_parser};

use std::ops::{Add, AddAssign, BitAnd, BitOr, Mul, Sub, Shl, ShlAssign, Shr};
use std::io;

/// An interface for parsing through some type of producer to a byte stream.
///
/// External parsers get passed in and consumes the bytes held internally
/// and outputs the `Result` of that parser.
pub trait StreamProducer {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T, ErrorKind>;
}

/// An abstraction trait for keeping different sized integers.
pub trait Sample: PartialEq + Eq + Sized + Clone + Copy +
                  Add<Output = Self> + AddAssign +
                  BitAnd<Self, Output = Self> + BitOr<Self, Output = Self> +
                  Mul<Output = Self> + Shl<u32, Output = Self>  +
                  ShlAssign<u32> + Shr<u32, Output = Self> +
                  Shr<i8, Output = Self> + Shr<i32, Output = Self> +
                  Sub<Output = Self> {
  /// The normal size for the current a `Sample`.
  type Normal;

  /// The size, in bits, for the `Sample::Normal`.
  fn size() -> usize;

  /// The size, in bits, for the `Sample`.
  fn size_extended() -> usize;

  /// Convert the extended `Sample` to the normal.
  fn to_normal(sample: Self) -> Option<Self::Normal>;

  /// Convert an i8 into a `Sample`.
  fn from_i8(sample: i8) -> Self;

  /// Convert an i16 into a `Sample`.
  fn from_i16(sample: i16) -> Self;

  /// Convert an i32 into a `Sample`.
  ///
  /// With `Sample` sometimes being smaller than a i32, there is a chance
  /// for this function to return an incorrect number. So when the number is
  /// larger of smaller than the current `Sample`, it returns `None`
  /// otherwise `Some(sample)`.
  fn from_i32(sample: i32) -> Option<Self>;

  /// Convert an i32 into a `Sample`.
  fn from_i32_lossy(sample: i32) -> Self;
}

/// A trait for defining the size of a sample.
pub trait SampleSize {
  /// The internal integer size used with `Stream::iter`.
  ///
  /// Rather than making the user of `Stream::iter` will remember what the
  /// extended size of the output sample will be, this is to map to the
  /// value based on the current integer size used.
  type Extended: Sample;
}

impl SampleSize for i8 {
  type Extended = i16;
}

impl SampleSize for i16 {
  type Extended = i32;
}

impl SampleSize for i32 {
  type Extended = i64;
}

pub trait WriteExtension: io::Write {
  fn write_u8(&mut self, number: u8) -> io::Result<()>;

  fn write_be_u16(&mut self, number: u16) -> io::Result<()>;
  fn write_le_u16(&mut self, number: u16) -> io::Result<()>;

  fn write_be_u24(&mut self, number: u32) -> io::Result<()>;
  fn write_le_u24(&mut self, number: u32) -> io::Result<()>;

  fn write_be_u32(&mut self, number: u32) -> io::Result<()>;
  fn write_le_u32(&mut self, number: u32) -> io::Result<()>;

  fn write_be_u64(&mut self, number: u64) -> io::Result<()>;
  fn write_le_u64(&mut self, number: u64) -> io::Result<()>;
}

impl<Write> WriteExtension for Write where Write: io::Write {
  fn write_u8(&mut self, number: u8) -> io::Result<()> {
    self.write_all(&[number])
  }

  fn write_be_u16(&mut self, number: u16) -> io::Result<()> {
    let mut buffer = [0; 2];

    buffer[0] = (number >> 8) as u8;
    buffer[1] = number as u8;

    self.write_all(&buffer)
  }

  fn write_le_u16(&mut self, number: u16) -> io::Result<()> {
    let mut buffer = [0; 2];

    buffer[2] = number as u8;
    buffer[1] = (number >> 8) as u8;

    self.write_all(&buffer)
  }

  fn write_be_u24(&mut self, number: u32) -> io::Result<()> {
    let mut buffer = [0; 3];

    buffer[0] = (number >> 16) as u8;
    buffer[1] = (number >> 8) as u8;
    buffer[2] = number as u8;

    self.write_all(&buffer)
  }

  fn write_le_u24(&mut self, number: u32) -> io::Result<()> {
    let mut buffer = [0; 3];

    buffer[0] = number as u8;
    buffer[1] = (number >> 8) as u8;
    buffer[2] = (number >> 16) as u8;

    self.write_all(&buffer)
  }
}

// Convert one to four byte slices into an unsigned 32-bit number.
//
// NOTE: This assumes big-endian since most numbers in the FLAC binary are
// that endianness.
#[inline]
pub fn to_u32(bytes: &[u8]) -> u32 {
  let length = bytes.len();

  debug_assert!(length <= 4);

  (0..length).fold(0, |result, i|
    result + ((bytes[i] as u32) << ((length - 1 - i) * 8))
  )
}

// Extends a signed value of a specific bit size to a larger bit size.
//
// NOTE: This assumes that the larger bit size will be 32 bit since that is
// the largest sample size supported in FLAC.
pub fn extend_sign(value: u32, bit_count: usize) -> i32 {
  if bit_count >= 32 || value < (1 << (bit_count - 1)) {
    value as i32
  } else {
    (value as i32).wrapping_sub(1 << bit_count)
  }
}

// Bit shifted version for two to the power of a given exponent.
#[inline]
pub fn power_of_two(exponent: u32) -> u32 {
  debug_assert!(exponent <= 31);

  1 << exponent
}

#[derive(PartialEq, Eq)]
enum ParserState {
  Header,
  StreamInfo,
  Metadata
}

fn parser<'a>(input: &'a [u8], state: &mut ParserState)
              -> IResult<&'a [u8], Metadata, ErrorKind> {
  let mut slice = input;
  let error     = nom::Err::Code(nom::ErrorKind::Custom(ErrorKind::Unknown));

  if *state == ParserState::Header {
    let (i, _) = try_parser! {
      to_custom_error!(slice, tag!("fLaC"), HeaderParser)
    };

    slice  = i;
    *state = ParserState::StreamInfo;
  }

  match *state {
    ParserState::StreamInfo => {
      let (i, block) = try_parse!(slice, metadata_parser);

      if block.is_stream_info() {
        *state = ParserState::Metadata;

        IResult::Done(i, block)
      } else {
        IResult::Error(error)
      }
    }
    ParserState::Metadata   => metadata_parser(slice),
    _                       => IResult::Error(error),
  }
}

pub fn many_metadata<S, F>(stream: &mut S, mut f: F) -> Result<(), ErrorKind>
 where S: StreamProducer,
       F: FnMut(Metadata) {
  let mut state  = ParserState::Header;
  let mut result = Ok(());

  loop {
    match stream.parse(|i| parser(i, &mut state)) {
      Ok(block)                => {
        let is_last = block.is_last();

        f(block);

        if is_last {
          break;
        }
      }
      Err(ErrorKind::Continue) => continue,
      Err(e)                   => {
        result = Err(e);

        break;
      }
    }
  }

  result
}

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

  #[test]
  #[should_panic]
  fn test_panic_to_u32() {
    to_u32(&[0x00, 0x00, 0x00, 0x00, 0x00]);
  }

  #[test]
  fn test_to_u32() {
    let bytes = [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef];

    assert_eq!(to_u32(&bytes[0..1]), 0x00000001);
    assert_eq!(to_u32(&bytes[3..5]), 0x00006789);
    assert_eq!(to_u32(&bytes[1..4]), 0x00234567);
    assert_eq!(to_u32(&bytes[4..]), 0x89abcdef);
  }

  #[test]
  fn test_extend_sign() {
    assert_eq!(extend_sign(32, 6), -32);
    assert_eq!(extend_sign(31, 6), 31);
    assert_eq!(extend_sign(128, 8), -128);
    assert_eq!(extend_sign(127, 8), 127);

    assert_eq!(extend_sign(2048, 12), -2048);
    assert_eq!(extend_sign(2047, 12), 2047);
    assert_eq!(extend_sign(32768, 16), -32768);
    assert_eq!(extend_sign(32767, 16), 32767);

    assert_eq!(extend_sign(8388608, 24), -8388608);
    assert_eq!(extend_sign(8388607, 24), 8388607);
    assert_eq!(extend_sign(2147483648, 32), -2147483648);
    assert_eq!(extend_sign(2147483647, 32), 2147483647);
  }

  #[test]
  #[should_panic]
  fn test_panic_power_of_two() {
    power_of_two(32);
  }

  #[test]
  fn test_power_of_two() {
    assert_eq!(power_of_two(0), 1);
    assert_eq!(power_of_two(1), 2);
    assert_eq!(power_of_two(2), 4);
    assert_eq!(power_of_two(10), 1024);
    assert_eq!(power_of_two(31), 2147483648);
  }
}
<|endoftext|>