<commit_before>use std::collections::btree_map::Entry;
use std::collections::{BTreeMap, BTreeSet};
use std::env;
use std::ffi::OsString;
use std::fs::{self, File};
use std::io::prelude::*;
use std::path::{Path, PathBuf};

use toml;

use core::{SourceId, Source, Package, Registry, Dependency, PackageIdSpec};
use core::PackageId;
use ops::{self, CompileFilter};
use sources::{GitSource, PathSource, RegistrySource};
use util::{CargoResult, ChainError, Config, human, internal};

#[derive(RustcDecodable, RustcEncodable)]
enum CrateListing {
    V1(CrateListingV1),
}

#[derive(RustcDecodable, RustcEncodable)]
struct CrateListingV1 {
    v1: BTreeMap<PackageId, BTreeSet<String>>,
}

struct Transaction {
    bins: Vec<PathBuf>,
}

impl Drop for Transaction {
    fn drop(&mut self) {
        for bin in self.bins.iter() {
            let _ = fs::remove_file(bin);
        }
    }
}

pub fn install(root: Option<&str>,
               krate: Option<&str>,
               source_id: &SourceId,
               vers: Option<&str>,
               opts: &ops::CompileOptions) -> CargoResult<()> {
    let config = opts.config;
    let root = try!(resolve_root(root, config));
    let (pkg, source) = if source_id.is_git() {
        try!(select_pkg(GitSource::new(source_id, config), source_id,
                        krate, vers, &mut |git| git.read_packages()))
    } else if source_id.is_path() {
        let path = source_id.url().to_file_path().ok()
                            .expect("path sources must have a valid path");
        try!(select_pkg(PathSource::new(&path, source_id, config),
                        source_id, krate, vers,
                        &mut |path| path.read_packages()))
    } else {
        try!(select_pkg(RegistrySource::new(source_id, config),
                        source_id, krate, vers,
                        &mut |_| Err(human("must specify a crate to install from \
                                            crates.io, or use --path or --git to \
                                            specify alternate source"))))
    };

    let mut list = try!(read_crate_list(&root));
    let dst = root.join("bin");
    try!(check_overwrites(&dst, &pkg, &opts.filter, &list));

    let target_dir = config.cwd().join("target-install");
    config.set_target_dir(&target_dir);
    let compile = try!(ops::compile_pkg(&pkg, Some(source), opts).chain_error(|| {
        human(format!("failed to compile `{}`, intermediate artifacts can be \
                       found at `{}`", pkg, target_dir.display()))
    }));

    let mut t = Transaction { bins: Vec::new() };
    try!(fs::create_dir_all(&dst));
    for bin in compile.binaries.iter() {
        let dst = dst.join(bin.file_name().unwrap());
        try!(config.shell().status("Installing", dst.display()));
        try!(fs::copy(&bin, &dst).chain_error(|| {
            human(format!("failed to copy `{}` to `{}`", bin.display(),
                          dst.display()))
        }));
        t.bins.push(dst);
    }
    try!(fs::remove_dir_all(&target_dir));

    list.v1.entry(pkg.package_id().clone()).or_insert_with(|| {
        BTreeSet::new()
    }).extend(t.bins.iter().map(|t| {
        t.file_name().unwrap().to_string_lossy().into_owned()
    }));
    try!(write_crate_list(&root, list));

    t.bins.truncate(0);

    // Print a warning that if this directory isn't in PATH that they won't be
    // able to run these commands.
    let path = env::var_os("PATH").unwrap_or(OsString::new());
    for path in env::split_paths(&path) {
        if path == dst {
            return Ok(())
        }
    }

    try!(config.shell().warn(&format!("be sure to add `{}` to your PATH to be \
                                       able to run the installed binaries",
                                      dst.display())));
    Ok(())
}

fn select_pkg<'a, T>(mut source: T,
                     source_id: &SourceId,
                     name: Option<&str>,
                     vers: Option<&str>,
                     list_all: &mut FnMut(&mut T) -> CargoResult<Vec<Package>>)
                     -> CargoResult<(Package, Box<Source + 'a>)>
    where T: Source + 'a
{
    try!(source.update());
    match name {
        Some(name) => {
            let dep = try!(Dependency::parse(name, vers, source_id));
            let deps = try!(source.query(&dep));
            match deps.iter().map(|p| p.package_id()).max() {
                Some(pkgid) => {
                    try!(source.download(&[pkgid.clone()]));
                    Ok((try!(source.get(&[pkgid.clone()])).remove(0),
                        Box::new(source)))
                }
                None => {
                    let vers_info = vers.map(|v| format!(" with version `{}`", v))
                                        .unwrap_or(String::new());
                    Err(human(format!("could not find `{}` in `{}`{}", name,
                                      source_id, vers_info)))
                }
            }
        }
        None => {
            let candidates = try!(list_all(&mut source));
            let binaries = candidates.iter().filter(|cand| {
                cand.targets().iter().filter(|t| t.is_bin()).count() > 0
            });
            let examples = candidates.iter().filter(|cand| {
                cand.targets().iter().filter(|t| t.is_example()).count() > 0
            });
            let pkg = match try!(one(binaries, |v| multi_err("binaries", v))) {
                Some(p) => p,
                None => {
                    match try!(one(examples, |v| multi_err("examples", v))) {
                        Some(p) => p,
                        None => bail!("no packages found with binaries or \
                                       examples"),
                    }
                }
            };
            return Ok((pkg.clone(), Box::new(source)));

            #[allow(deprecated)] // connect => join in 1.3
            fn multi_err(kind: &str, mut pkgs: Vec<&Package>) -> String {
                pkgs.sort_by(|a, b| a.name().cmp(b.name()));
                format!("multiple packages with {} found: {}", kind,
                        pkgs.iter().map(|p| p.name()).collect::<Vec<_>>()
                            .connect(", "))
            }
        }
    }
}

fn one<I, F>(mut i: I, f: F) -> CargoResult<Option<I::Item>>
    where I: Iterator,
          F: FnOnce(Vec<I::Item>) -> String
{
    match (i.next(), i.next()) {
        (Some(i1), Some(i2)) => {
            let mut v = vec![i1, i2];
            v.extend(i);
            Err(human(f(v)))
        }
        (Some(i), None) => Ok(Some(i)),
        (None, _) => Ok(None)
    }
}

fn check_overwrites(dst: &Path,
                    pkg: &Package,
                    filter: &ops::CompileFilter,
                    prev: &CrateListingV1) -> CargoResult<()> {
    let check = |name| {
        let name = format!("{}{}", name, env::consts::EXE_SUFFIX);
        if fs::metadata(dst.join(&name)).is_err() {
            return Ok(())
        }
        let mut msg = format!("binary `{}` already exists in destination", name);
        if let Some((p, _)) = prev.v1.iter().find(|&(_, v)| v.contains(&name)) {
            msg.push_str(&format!(" as part of `{}`", p));
        }
        Err(human(msg))
    };
    match *filter {
        CompileFilter::Everything => {
            // If explicit --bin or --example flags were passed then those'll
            // get checked during cargo_compile, we only care about the "build
            // everything" case here
            if pkg.targets().iter().filter(|t| t.is_bin()).next().is_none() {
                bail!("specified package has no binaries")
            }

            for target in pkg.targets().iter().filter(|t| t.is_bin()) {
                try!(check(target.name()));
            }
        }
        CompileFilter::Only { bins, examples, .. } => {
            for bin in bins.iter().chain(examples) {
                try!(check(bin));
            }
        }
    }
    Ok(())
}

fn read_crate_list(path: &Path) -> CargoResult<CrateListingV1> {
    let metadata = path.join(".crates.toml");
    let mut f = match File::open(&metadata) {
        Ok(f) => f,
        Err(..) => return Ok(CrateListingV1 { v1: BTreeMap::new() }),
    };
    (|| -> CargoResult<_> {
        let mut contents = String::new();
        try!(f.read_to_string(&mut contents));
        let listing = try!(toml::decode_str(&contents).chain_error(|| {
            internal("invalid TOML found for metadata")
        }));
        match listing {
            CrateListing::V1(v1) => Ok(v1),
        }
    }).chain_error(|| {
        human(format!("failed to parse crate metadata at `{}`",
                      metadata.display()))
    })
}

fn write_crate_list(path: &Path, listing: CrateListingV1) -> CargoResult<()> {
    let metadata = path.join(".crates.toml");
    (|| -> CargoResult<_> {
        let mut f = try!(File::create(&metadata));
        let data = toml::encode_str::<CrateListing>(&CrateListing::V1(listing));
        try!(f.write_all(data.as_bytes()));
        Ok(())
    }).chain_error(|| {
        human(format!("failed to write crate metadata at `{}`",
                      metadata.display()))
    })
}

pub fn install_list(dst: Option<&str>, config: &Config) -> CargoResult<()> {
    let dst = try!(resolve_root(dst, config));
    let list = try!(read_crate_list(&dst));
    let mut shell = config.shell();
    let out = shell.out();
    for (k, v) in list.v1.iter() {
        try!(writeln!(out, "{}:", k));
        for bin in v {
            try!(writeln!(out, "    {}", bin));
        }
    }
    Ok(())
}

pub fn uninstall(root: Option<&str>,
                 spec: &str,
                 bins: &[String],
                 config: &Config) -> CargoResult<()> {
    let root = try!(resolve_root(root, config));
    let mut metadata = try!(read_crate_list(&root));
    let mut to_remove = Vec::new();
    {
        let result = try!(PackageIdSpec::query_str(spec, metadata.v1.keys()))
                                        .clone();
        let mut installed = match metadata.v1.entry(result.clone()) {
            Entry::Occupied(e) => e,
            Entry::Vacant(..) => panic!("entry not found: {}", result),
        };
        let dst = root.join("bin");
        for bin in installed.get() {
            let bin = dst.join(bin);
            if fs::metadata(&bin).is_err() {
                bail!("corrupt metadata, `{}` does not exist when it should",
                      bin.display())
            }
        }

        let bins = bins.iter().map(|s| {
            if s.ends_with(env::consts::EXE_SUFFIX) {
                s.to_string()
            } else {
                format!("{}{}", s, env::consts::EXE_SUFFIX)
            }
        }).collect::<Vec<_>>();

        for bin in bins.iter() {
            if !installed.get().contains(bin) {
                bail!("binary `{}` not installed as part of `{}`", bin, result)
            }
        }

        if bins.len() == 0 {
            to_remove.extend(installed.get().iter().map(|b| dst.join(b)));
            installed.get_mut().clear();
        } else {
            for bin in bins.iter() {
                to_remove.push(dst.join(bin));
                installed.get_mut().remove(bin);
            }
        }
        if installed.get().len() == 0 {
            installed.remove();
        }
    }
    try!(write_crate_list(&root, metadata));
    for bin in to_remove {
        try!(config.shell().status("Removing", bin.display()));
        try!(fs::remove_file(bin));
    }

    Ok(())
}

fn resolve_root(flag: Option<&str>, config: &Config) -> CargoResult<PathBuf> {
    let config_root = try!(config.get_string("install.root"));
    Ok(flag.map(PathBuf::from).or_else(|| {
        env::var_os("CARGO_INSTALL_ROOT").map(PathBuf::from)
    }).or_else(|| {
        config_root.clone().map(|(v, _)| PathBuf::from(v))
    }).unwrap_or_else(|| {
        config.home().to_owned()
    }))
}
<commit_msg>Only treat missing crate lists as empty<commit_after>use std::collections::btree_map::Entry;
use std::collections::{BTreeMap, BTreeSet};
use std::env;
use std::ffi::OsString;
use std::fs::{self, File};
use std::io::prelude::*;
use std::io;
use std::path::{Path, PathBuf};

use toml;

use core::{SourceId, Source, Package, Registry, Dependency, PackageIdSpec};
use core::PackageId;
use ops::{self, CompileFilter};
use sources::{GitSource, PathSource, RegistrySource};
use util::{CargoResult, ChainError, Config, human, internal};

#[derive(RustcDecodable, RustcEncodable)]
enum CrateListing {
    V1(CrateListingV1),
}

#[derive(RustcDecodable, RustcEncodable)]
struct CrateListingV1 {
    v1: BTreeMap<PackageId, BTreeSet<String>>,
}

struct Transaction {
    bins: Vec<PathBuf>,
}

impl Drop for Transaction {
    fn drop(&mut self) {
        for bin in self.bins.iter() {
            let _ = fs::remove_file(bin);
        }
    }
}

pub fn install(root: Option<&str>,
               krate: Option<&str>,
               source_id: &SourceId,
               vers: Option<&str>,
               opts: &ops::CompileOptions) -> CargoResult<()> {
    let config = opts.config;
    let root = try!(resolve_root(root, config));
    let (pkg, source) = if source_id.is_git() {
        try!(select_pkg(GitSource::new(source_id, config), source_id,
                        krate, vers, &mut |git| git.read_packages()))
    } else if source_id.is_path() {
        let path = source_id.url().to_file_path().ok()
                            .expect("path sources must have a valid path");
        try!(select_pkg(PathSource::new(&path, source_id, config),
                        source_id, krate, vers,
                        &mut |path| path.read_packages()))
    } else {
        try!(select_pkg(RegistrySource::new(source_id, config),
                        source_id, krate, vers,
                        &mut |_| Err(human("must specify a crate to install from \
                                            crates.io, or use --path or --git to \
                                            specify alternate source"))))
    };

    let mut list = try!(read_crate_list(&root));
    let dst = root.join("bin");
    try!(check_overwrites(&dst, &pkg, &opts.filter, &list));

    let target_dir = config.cwd().join("target-install");
    config.set_target_dir(&target_dir);
    let compile = try!(ops::compile_pkg(&pkg, Some(source), opts).chain_error(|| {
        human(format!("failed to compile `{}`, intermediate artifacts can be \
                       found at `{}`", pkg, target_dir.display()))
    }));

    let mut t = Transaction { bins: Vec::new() };
    try!(fs::create_dir_all(&dst));
    for bin in compile.binaries.iter() {
        let dst = dst.join(bin.file_name().unwrap());
        try!(config.shell().status("Installing", dst.display()));
        try!(fs::copy(&bin, &dst).chain_error(|| {
            human(format!("failed to copy `{}` to `{}`", bin.display(),
                          dst.display()))
        }));
        t.bins.push(dst);
    }
    try!(fs::remove_dir_all(&target_dir));

    list.v1.entry(pkg.package_id().clone()).or_insert_with(|| {
        BTreeSet::new()
    }).extend(t.bins.iter().map(|t| {
        t.file_name().unwrap().to_string_lossy().into_owned()
    }));
    try!(write_crate_list(&root, list));

    t.bins.truncate(0);

    // Print a warning that if this directory isn't in PATH that they won't be
    // able to run these commands.
    let path = env::var_os("PATH").unwrap_or(OsString::new());
    for path in env::split_paths(&path) {
        if path == dst {
            return Ok(())
        }
    }

    try!(config.shell().warn(&format!("be sure to add `{}` to your PATH to be \
                                       able to run the installed binaries",
                                      dst.display())));
    Ok(())
}

fn select_pkg<'a, T>(mut source: T,
                     source_id: &SourceId,
                     name: Option<&str>,
                     vers: Option<&str>,
                     list_all: &mut FnMut(&mut T) -> CargoResult<Vec<Package>>)
                     -> CargoResult<(Package, Box<Source + 'a>)>
    where T: Source + 'a
{
    try!(source.update());
    match name {
        Some(name) => {
            let dep = try!(Dependency::parse(name, vers, source_id));
            let deps = try!(source.query(&dep));
            match deps.iter().map(|p| p.package_id()).max() {
                Some(pkgid) => {
                    try!(source.download(&[pkgid.clone()]));
                    Ok((try!(source.get(&[pkgid.clone()])).remove(0),
                        Box::new(source)))
                }
                None => {
                    let vers_info = vers.map(|v| format!(" with version `{}`", v))
                                        .unwrap_or(String::new());
                    Err(human(format!("could not find `{}` in `{}`{}", name,
                                      source_id, vers_info)))
                }
            }
        }
        None => {
            let candidates = try!(list_all(&mut source));
            let binaries = candidates.iter().filter(|cand| {
                cand.targets().iter().filter(|t| t.is_bin()).count() > 0
            });
            let examples = candidates.iter().filter(|cand| {
                cand.targets().iter().filter(|t| t.is_example()).count() > 0
            });
            let pkg = match try!(one(binaries, |v| multi_err("binaries", v))) {
                Some(p) => p,
                None => {
                    match try!(one(examples, |v| multi_err("examples", v))) {
                        Some(p) => p,
                        None => bail!("no packages found with binaries or \
                                       examples"),
                    }
                }
            };
            return Ok((pkg.clone(), Box::new(source)));

            #[allow(deprecated)] // connect => join in 1.3
            fn multi_err(kind: &str, mut pkgs: Vec<&Package>) -> String {
                pkgs.sort_by(|a, b| a.name().cmp(b.name()));
                format!("multiple packages with {} found: {}", kind,
                        pkgs.iter().map(|p| p.name()).collect::<Vec<_>>()
                            .connect(", "))
            }
        }
    }
}

fn one<I, F>(mut i: I, f: F) -> CargoResult<Option<I::Item>>
    where I: Iterator,
          F: FnOnce(Vec<I::Item>) -> String
{
    match (i.next(), i.next()) {
        (Some(i1), Some(i2)) => {
            let mut v = vec![i1, i2];
            v.extend(i);
            Err(human(f(v)))
        }
        (Some(i), None) => Ok(Some(i)),
        (None, _) => Ok(None)
    }
}

fn check_overwrites(dst: &Path,
                    pkg: &Package,
                    filter: &ops::CompileFilter,
                    prev: &CrateListingV1) -> CargoResult<()> {
    let check = |name| {
        let name = format!("{}{}", name, env::consts::EXE_SUFFIX);
        if fs::metadata(dst.join(&name)).is_err() {
            return Ok(())
        }
        let mut msg = format!("binary `{}` already exists in destination", name);
        if let Some((p, _)) = prev.v1.iter().find(|&(_, v)| v.contains(&name)) {
            msg.push_str(&format!(" as part of `{}`", p));
        }
        Err(human(msg))
    };
    match *filter {
        CompileFilter::Everything => {
            // If explicit --bin or --example flags were passed then those'll
            // get checked during cargo_compile, we only care about the "build
            // everything" case here
            if pkg.targets().iter().filter(|t| t.is_bin()).next().is_none() {
                bail!("specified package has no binaries")
            }

            for target in pkg.targets().iter().filter(|t| t.is_bin()) {
                try!(check(target.name()));
            }
        }
        CompileFilter::Only { bins, examples, .. } => {
            for bin in bins.iter().chain(examples) {
                try!(check(bin));
            }
        }
    }
    Ok(())
}

fn read_crate_list(path: &Path) -> CargoResult<CrateListingV1> {
    let metadata = path.join(".crates.toml");
    let mut f = match File::open(&metadata) {
        Ok(f) => f,
        Err(e) => {
            if e.kind() == io::ErrorKind::NotFound {
                return Ok(CrateListingV1 { v1: BTreeMap::new() });
            }
            return Err(e).chain_error(|| {
                human(format!("failed to open crate metadata at `{}`",
                              metadata.display()))
            });
        }
    };
    (|| -> CargoResult<_> {
        let mut contents = String::new();
        try!(f.read_to_string(&mut contents));
        let listing = try!(toml::decode_str(&contents).chain_error(|| {
            internal("invalid TOML found for metadata")
        }));
        match listing {
            CrateListing::V1(v1) => Ok(v1),
        }
    }).chain_error(|| {
        human(format!("failed to parse crate metadata at `{}`",
                      metadata.display()))
    })
}

fn write_crate_list(path: &Path, listing: CrateListingV1) -> CargoResult<()> {
    let metadata = path.join(".crates.toml");
    (|| -> CargoResult<_> {
        let mut f = try!(File::create(&metadata));
        let data = toml::encode_str::<CrateListing>(&CrateListing::V1(listing));
        try!(f.write_all(data.as_bytes()));
        Ok(())
    }).chain_error(|| {
        human(format!("failed to write crate metadata at `{}`",
                      metadata.display()))
    })
}

pub fn install_list(dst: Option<&str>, config: &Config) -> CargoResult<()> {
    let dst = try!(resolve_root(dst, config));
    let list = try!(read_crate_list(&dst));
    let mut shell = config.shell();
    let out = shell.out();
    for (k, v) in list.v1.iter() {
        try!(writeln!(out, "{}:", k));
        for bin in v {
            try!(writeln!(out, "    {}", bin));
        }
    }
    Ok(())
}

pub fn uninstall(root: Option<&str>,
                 spec: &str,
                 bins: &[String],
                 config: &Config) -> CargoResult<()> {
    let root = try!(resolve_root(root, config));
    let mut metadata = try!(read_crate_list(&root));
    let mut to_remove = Vec::new();
    {
        let result = try!(PackageIdSpec::query_str(spec, metadata.v1.keys()))
                                        .clone();
        let mut installed = match metadata.v1.entry(result.clone()) {
            Entry::Occupied(e) => e,
            Entry::Vacant(..) => panic!("entry not found: {}", result),
        };
        let dst = root.join("bin");
        for bin in installed.get() {
            let bin = dst.join(bin);
            if fs::metadata(&bin).is_err() {
                bail!("corrupt metadata, `{}` does not exist when it should",
                      bin.display())
            }
        }

        let bins = bins.iter().map(|s| {
            if s.ends_with(env::consts::EXE_SUFFIX) {
                s.to_string()
            } else {
                format!("{}{}", s, env::consts::EXE_SUFFIX)
            }
        }).collect::<Vec<_>>();

        for bin in bins.iter() {
            if !installed.get().contains(bin) {
                bail!("binary `{}` not installed as part of `{}`", bin, result)
            }
        }

        if bins.len() == 0 {
            to_remove.extend(installed.get().iter().map(|b| dst.join(b)));
            installed.get_mut().clear();
        } else {
            for bin in bins.iter() {
                to_remove.push(dst.join(bin));
                installed.get_mut().remove(bin);
            }
        }
        if installed.get().len() == 0 {
            installed.remove();
        }
    }
    try!(write_crate_list(&root, metadata));
    for bin in to_remove {
        try!(config.shell().status("Removing", bin.display()));
        try!(fs::remove_file(bin));
    }

    Ok(())
}

fn resolve_root(flag: Option<&str>, config: &Config) -> CargoResult<PathBuf> {
    let config_root = try!(config.get_string("install.root"));
    Ok(flag.map(PathBuf::from).or_else(|| {
        env::var_os("CARGO_INSTALL_ROOT").map(PathBuf::from)
    }).or_else(|| {
        config_root.clone().map(|(v, _)| PathBuf::from(v))
    }).unwrap_or_else(|| {
        config.home().to_owned()
    }))
}
<|endoftext|>


<commit_before>
//! Append another stream.

use event;
use modifier;
use transform;
use builder;

/// Appends another stream after the current.
#[derive(Debug)]
pub struct Append<S, N> {
    stream: modifier::Combine<S, N>,
}

impl<S, N> Append<S, N>
where
    S: event::Stream,
    N: event::Stream,
{
    pub(crate) fn new(stream: S, new_stream: N) -> Append<S, N> {
        Append {
            stream: modifier::Combine::new(stream, new_stream),
        }
    }
}

impl<S, N> event::Stream for Append<S, N>
where
    S: event::Stream,
    N: event::Stream,
{
    fn next_event(&mut self) -> event::StreamResult {
        self.stream.next_event()
    }
}

/// Appends another stream to the contents of the current element.
pub struct AppendContent<S, N>
where
    S: event::ElementStream,
    N: event::Stream,
{
    stream: modifier::select::SelectContent<S, Builder<N>>,
}

impl<S, N> AppendContent<S, N>
where
    S: event::ElementStream,
    N: event::Stream,
{
    pub(crate) fn new(stream: S, new_stream: N) -> AppendContent<S, N> {
        AppendContent {
            stream: modifier::select::SelectContent::new(stream, Builder { new_stream }),
        }
    }
}

impl<S, N> event::ElementStream for AppendContent<S, N>
where
    S: event::ElementStream,
    N: event::Stream,
{}

impl<S, N> event::Stream for AppendContent<S, N>
where
    S: event::ElementStream,
    N: event::Stream,
{
    fn next_event(&mut self) -> event::StreamResult {
        self.stream.next_event()
    }
}

struct Builder<N> {
    new_stream: N,
}

impl<S, N> builder::BuildOnce<modifier::select::CurrentContent<S>> for Builder<N>
where
    S: event::ElementStream,
    N: event::Stream,
{
    type Stream = Append<modifier::select::CurrentContent<S>, N>;

    fn build_once(self, stream: transform::Api<modifier::select::CurrentContent<S>>)
    -> transform::Api<Self::Stream> {
        transform::Api::pack(Append::new(stream.unpack(), self.new_stream))
    }
}

#[cfg(test)]
mod tests {

    test_group!(append_contents:
        "append to existing" => transform_test!(
            "23<c>42</c>93",
            "23<c>42foo</c>93",
            |html| html.select("c", |html| html.append_contents("foo")),
        ),
        "append to empty" => transform_test!(
            "23<c></c>93",
            "23<c>foo</c>93",
            |html| html.select("c", |html| html.append_contents("foo")),
        ),
    );

    test_group!(append:
        "append" => transform_test!(
            "23<c>42</c>93",
            "23<c>42</c>foo93",
            |html| html.select("c", |html| html.append("foo")),
        ),
    );
}

<commit_msg>added element stream tests to append<commit_after>
//! Append another stream.

use event;
use modifier;
use transform;
use builder;

/// Appends another stream after the current.
#[derive(Debug)]
pub struct Append<S, N> {
    stream: modifier::Combine<S, N>,
}

impl<S, N> Append<S, N>
where
    S: event::Stream,
    N: event::Stream,
{
    pub(crate) fn new(stream: S, new_stream: N) -> Append<S, N> {
        Append {
            stream: modifier::Combine::new(stream, new_stream),
        }
    }
}

impl<S, N> event::Stream for Append<S, N>
where
    S: event::Stream,
    N: event::Stream,
{
    fn next_event(&mut self) -> event::StreamResult {
        self.stream.next_event()
    }
}

/// Appends another stream to the contents of the current element.
pub struct AppendContent<S, N>
where
    S: event::ElementStream,
    N: event::Stream,
{
    stream: modifier::select::SelectContent<S, Builder<N>>,
}

impl<S, N> AppendContent<S, N>
where
    S: event::ElementStream,
    N: event::Stream,
{
    pub(crate) fn new(stream: S, new_stream: N) -> AppendContent<S, N> {
        AppendContent {
            stream: modifier::select::SelectContent::new(stream, Builder { new_stream }),
        }
    }
}

impl<S, N> event::ElementStream for AppendContent<S, N>
where
    S: event::ElementStream,
    N: event::Stream,
{}

impl<S, N> event::Stream for AppendContent<S, N>
where
    S: event::ElementStream,
    N: event::Stream,
{
    fn next_event(&mut self) -> event::StreamResult {
        self.stream.next_event()
    }
}

struct Builder<N> {
    new_stream: N,
}

impl<S, N> builder::BuildOnce<modifier::select::CurrentContent<S>> for Builder<N>
where
    S: event::ElementStream,
    N: event::Stream,
{
    type Stream = Append<modifier::select::CurrentContent<S>, N>;

    fn build_once(self, stream: transform::Api<modifier::select::CurrentContent<S>>)
    -> transform::Api<Self::Stream> {
        transform::Api::pack(Append::new(stream.unpack(), self.new_stream))
    }
}

#[cfg(test)]
mod tests {

    test_group!(append_contents:
        "append to existing" => transform_test!(
            "23<c>42</c>93",
            "23<c>42foo</c>93",
            |html| html.select("c", |html| html.append_contents("foo")),
        ),
        "append to empty" => transform_test!(
            "23<c></c>93",
            "23<c>foo</c>93",
            |html| html.select("c", |html| html.append_contents("foo")),
        ),
        "ensure element stream" => transform_test!(
            "23<c></c>93",
            "23<c>foo</c>93",
            |html| html.select("c", |html| html.append_contents("foo").into_boxed_element()),
        ),
    );

    test_group!(append:
        "append" => transform_test!(
            "23<c>42</c>93",
            "23<c>42</c>foo93",
            |html| html.select("c", |html| html.append("foo")),
        ),
    );
}

<|endoftext|>


<commit_before>use std::path;
use std::process;

use clap::{App, Arg, SubCommand, AppSettings};

use config;
use models;
use web;
use utils;

macro_rules! clap_dispatch {
    // Transform `foo(_) => ()` into `foo(_,) => ()` such that it actually can be parsed.
    // Using a single rule is only possible in nightly.
    ($matches:expr; { $( $name:ident ( $matches_name:pat ) => $callback:expr ),* }) => (clap_dispatch!($matches; { $( $name($matches_name,) => $callback ),* }));

    ($matches:expr; { $( $name:ident ($matches_name:pat, $($arg_name:ident as $varname:ident),*) => $callback:expr ),* }) => {
        match $matches.subcommand_name() {
            $(
                Some(stringify!($name)) => {
                    let matches = $matches.subcommand_matches(stringify!($name)).unwrap();
                    $(
                        let $varname = matches.value_of(stringify!($arg_name)).unwrap();
                    )*
                    let $matches_name = matches;
                    $callback;
                }
            )*
            Some(x) => {
                panic!("Internal error: Command not covered: {}", x);
            },
            None => {
                println!("Subcommand required. See --help for help.");
                process::exit(1);
            }
        }
    };
}

pub fn main() {
    let matches =
        App::new("mysteryshack")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Markus Unterwaditzer & contributors")
        .about("A remoteStorage server.")
        .args_from_usage("-c, --config=[FILE] 'Use specified config file, defaults to ./config'")
        .setting(AppSettings::SubcommandRequired)
        .subcommand(SubCommand::with_name("serve")
                    .about("Start server"))
        .subcommand(SubCommand::with_name("user")
                    .about("User management")
                    .setting(AppSettings::SubcommandRequired)
                    .arg(Arg::with_name("USERNAME")
                         .help("The username to perform the operation with.")
                         .required(true)
                         .index(1))
                    .subcommand(SubCommand::with_name("create")
                                .about("Create a new user"))
                    .subcommand(SubCommand::with_name("setpass")
                                .about("Change password for user"))
                    .subcommand(SubCommand::with_name("delete")
                                .about("Delete a user"))
                    .subcommand(SubCommand::with_name("authorize")
                                .about("Create a OAuth token. This is mostly useful for development.")
                                .arg(Arg::with_name("CLIENT_ID").required(true).index(1))
                                .arg(Arg::with_name("SCOPE").required(true).index(2))))
        .get_matches();

    let config_path = path::Path::new(matches.value_of("config").unwrap_or("./config"));

    let config = match config::Config::read_file(config_path) {
        Ok(x) => x,
        Err(e) => {
            println!("Failed to parse config: {}", e);
            process::exit(1);
        }
    };

    clap_dispatch!(matches; {
        serve(_,) => web::run_server(config),  // FIXME: Bug in clap_dispatch: comma required
        user(user_matches, USERNAME as username) => clap_dispatch!(user_matches; {
            create(_,) => {
                let password_hash = models::PasswordHash::from_password(utils::double_prompt("Password for new user: "));

                match models::User::create(&config.data_path, username).map(|user| {
                    user.set_password_hash(password_hash)
                }) {
                    Ok(_) => (),
                    Err(e) => {
                        println!("Failed to create user {}: {}", username, e);
                        process::exit(1);
                    }
                };

                println!("Successfully created user {}", username);
            },
            setpass(_,) => {
                let user = match models::User::get(&config.data_path, username) {
                    Some(x) => x,
                    None => {
                        println!("User does not exist: {}", username);
                        process::exit(1);
                    }
                };

                let password_hash = models::PasswordHash::from_password(utils::double_prompt("New password: "));
                match user.set_password_hash(password_hash) {
                    Ok(_) => (),
                    Err(e) => {
                        println!("Failed to set password for user {}: {}", username, e);
                        process::exit(1);
                    }
                };

                println!("Changed password for user {}", username);
            },
            delete(_,) => {
                let user = match models::User::get(&config.data_path, username) {
                    Some(x) => x,
                    None => {
                        println!("User does not exist: {}", username);
                        process::exit(1);
                    }
                };

                println!("You are about to delete the user {} and ALL the user's user data. This
                         process is irreversible.", username);
                if !utils::prompt_confirm(format!("Do you want to delete the user {}?", username), false) {
                    println!("Aborted!");
                    process::exit(1);
                }

                match user.delete() {
                    Ok(_) => println!("Successfully deleted user {}.", username),
                    Err(e) => {
                        println!("Failed to delete user: {:?}", e);
                        process::exit(1);
                    }
                };
            },
            authorize(_, SCOPE as scope, CLIENT_ID as client_id) => {
                let user = match models::User::get(&config.data_path, username) {
                    Some(x) => x,
                    None => {
                        println!("User does not exist: {}", username);
                        process::exit(1);
                    }
                };

                let oauth_session = web::oauth::Session {
                    client_id: client_id.to_owned(),
                    permissions: match web::oauth::PermissionsMap::from_scope_string(scope) {
                        Some(x) => x,
                        None => {
                            println!("Invalid scope: {}", scope);
                            process::exit(1);
                        }
                    }
                };

                let (_, token) = models::Token::create(&user, oauth_session).unwrap();
                println!("{}", token.token(&user));
            }
        })
    });
}
<commit_msg>Init sodiumoxide<commit_after>use std::path;
use std::process;

use sodiumoxide;

use clap::{App, Arg, SubCommand, AppSettings};

use config;
use models;
use web;
use utils;

macro_rules! clap_dispatch {
    // Transform `foo(_) => ()` into `foo(_,) => ()` such that it actually can be parsed.
    // Using a single rule is only possible in nightly.
    ($matches:expr; { $( $name:ident ( $matches_name:pat ) => $callback:expr ),* }) => (clap_dispatch!($matches; { $( $name($matches_name,) => $callback ),* }));

    ($matches:expr; { $( $name:ident ($matches_name:pat, $($arg_name:ident as $varname:ident),*) => $callback:expr ),* }) => {
        match $matches.subcommand_name() {
            $(
                Some(stringify!($name)) => {
                    let matches = $matches.subcommand_matches(stringify!($name)).unwrap();
                    $(
                        let $varname = matches.value_of(stringify!($arg_name)).unwrap();
                    )*
                    let $matches_name = matches;
                    $callback;
                }
            )*
            Some(x) => {
                panic!("Internal error: Command not covered: {}", x);
            },
            None => {
                println!("Subcommand required. See --help for help.");
                process::exit(1);
            }
        }
    };
}

pub fn main() {
    let matches =
        App::new("mysteryshack")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Markus Unterwaditzer & contributors")
        .about("A remoteStorage server.")
        .args_from_usage("-c, --config=[FILE] 'Use specified config file, defaults to ./config'")
        .setting(AppSettings::SubcommandRequired)
        .subcommand(SubCommand::with_name("serve")
                    .about("Start server"))
        .subcommand(SubCommand::with_name("user")
                    .about("User management")
                    .setting(AppSettings::SubcommandRequired)
                    .arg(Arg::with_name("USERNAME")
                         .help("The username to perform the operation with.")
                         .required(true)
                         .index(1))
                    .subcommand(SubCommand::with_name("create")
                                .about("Create a new user"))
                    .subcommand(SubCommand::with_name("setpass")
                                .about("Change password for user"))
                    .subcommand(SubCommand::with_name("delete")
                                .about("Delete a user"))
                    .subcommand(SubCommand::with_name("authorize")
                                .about("Create a OAuth token. This is mostly useful for development.")
                                .arg(Arg::with_name("CLIENT_ID").required(true).index(1))
                                .arg(Arg::with_name("SCOPE").required(true).index(2))))
        .get_matches();

    assert!(sodiumoxide::init());

    let config_path = path::Path::new(matches.value_of("config").unwrap_or("./config"));

    let config = match config::Config::read_file(config_path) {
        Ok(x) => x,
        Err(e) => {
            println!("Failed to parse config: {}", e);
            process::exit(1);
        }
    };

    clap_dispatch!(matches; {
        serve(_,) => web::run_server(config),  // FIXME: Bug in clap_dispatch: comma required
        user(user_matches, USERNAME as username) => clap_dispatch!(user_matches; {
            create(_,) => {
                let password_hash = models::PasswordHash::from_password(utils::double_prompt("Password for new user: "));

                match models::User::create(&config.data_path, username).map(|user| {
                    user.set_password_hash(password_hash)
                }) {
                    Ok(_) => (),
                    Err(e) => {
                        println!("Failed to create user {}: {}", username, e);
                        process::exit(1);
                    }
                };

                println!("Successfully created user {}", username);
            },
            setpass(_,) => {
                let user = match models::User::get(&config.data_path, username) {
                    Some(x) => x,
                    None => {
                        println!("User does not exist: {}", username);
                        process::exit(1);
                    }
                };

                let password_hash = models::PasswordHash::from_password(utils::double_prompt("New password: "));
                match user.set_password_hash(password_hash) {
                    Ok(_) => (),
                    Err(e) => {
                        println!("Failed to set password for user {}: {}", username, e);
                        process::exit(1);
                    }
                };

                println!("Changed password for user {}", username);
            },
            delete(_,) => {
                let user = match models::User::get(&config.data_path, username) {
                    Some(x) => x,
                    None => {
                        println!("User does not exist: {}", username);
                        process::exit(1);
                    }
                };

                println!("You are about to delete the user {} and ALL the user's user data. This
                         process is irreversible.", username);
                if !utils::prompt_confirm(format!("Do you want to delete the user {}?", username), false) {
                    println!("Aborted!");
                    process::exit(1);
                }

                match user.delete() {
                    Ok(_) => println!("Successfully deleted user {}.", username),
                    Err(e) => {
                        println!("Failed to delete user: {:?}", e);
                        process::exit(1);
                    }
                };
            },
            authorize(_, SCOPE as scope, CLIENT_ID as client_id) => {
                let user = match models::User::get(&config.data_path, username) {
                    Some(x) => x,
                    None => {
                        println!("User does not exist: {}", username);
                        process::exit(1);
                    }
                };

                let oauth_session = web::oauth::Session {
                    client_id: client_id.to_owned(),
                    permissions: match web::oauth::PermissionsMap::from_scope_string(scope) {
                        Some(x) => x,
                        None => {
                            println!("Invalid scope: {}", scope);
                            process::exit(1);
                        }
                    }
                };

                let (_, token) = models::Token::create(&user, oauth_session).unwrap();
                println!("{}", token.token(&user));
            }
        })
    });
}
<|endoftext|>


<commit_before>// Copyright 2015 click2stream, Inc.
//
// 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.

//! PCAP network scanner definitions.

use std::fmt;
use std::ptr;
use std::slice;
use std::sync::mpsc;
use std::thread;

use std::error::Error;
use std::ffi::{CStr, CString};
use std::fmt::{Display, Formatter};
use std::sync::mpsc::Sender;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use bytes::{Bytes, BytesMut};

use libc::{c_char, c_int, c_void, size_t};

use crate::net::raw::ether::packet::EtherPacket;

type PacketCallback = unsafe extern "C" fn(
    opaque: *mut c_void,
    data: *const u8,
    data_length: size_t,
    packet_length: size_t,
);

extern "C" {
    fn pcap_wrapper__new(device: *const c_char) -> *mut c_void;
    fn pcap_wrapper__free(wrapper: *mut c_void);

    fn pcap_wrapper__get_last_error(wrapper: *const c_void) -> *const c_char;

    fn pcap_wrapper__set_filter(wrapper: *mut c_void, filter: *const c_char) -> c_int;
    fn pcap_wrapper__set_max_packet_length(wrapper: *mut c_void, max_packet_length: size_t);
    fn pcap_wrapper__set_read_timeout(wrapper: *mut c_void, read_timeout: u64);

    fn pcap_wrapper__open(wrapper: *mut c_void) -> c_int;

    fn pcap_wrapper__read_packet(
        wrapper: *mut c_void,
        callback: PacketCallback,
        opaque: *mut c_void,
    ) -> c_int;
    fn pcap_wrapper__write_packet(wrapper: *mut c_void, data: *const u8, size: size_t) -> c_int;
}

lazy_static! {
    /// PCAP context for synchronizing thread unsafe calls.
    static ref TC: Mutex<()> = Mutex::new(());
}

/// PCAP error.
#[derive(Debug, Clone)]
pub struct PcapError {
    msg: String,
}

impl PcapError {
    /// Create a new error with a given error message.
    pub fn new<T>(msg: T) -> Self
    where
        T: ToString,
    {
        Self {
            msg: msg.to_string(),
        }
    }
}

impl Display for PcapError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.write_str(&self.msg)
    }
}

impl Error for PcapError {}

/// Common result type for this module.
pub type Result<T> = std::result::Result<T, PcapError>;

/// Capture builder.
struct CaptureBuilder {
    inner: Capture,
}

impl CaptureBuilder {
    /// Create a new capture builder for a given device.
    fn new(device: &str) -> Self {
        let device = CString::new(device).unwrap();

        let ptr = unsafe { pcap_wrapper__new(device.as_ptr()) };

        if ptr.is_null() {
            panic!("Unable to allocate PCAP capture device");
        }

        let capture = Capture { ptr };

        Self { inner: capture }
    }

    /// Set maximum packet length (longer packets will be truncated).
    pub fn max_packet_length(self, max_packet_length: usize) -> Self {
        unsafe {
            pcap_wrapper__set_max_packet_length(self.inner.ptr, max_packet_length as _);
        }

        self
    }

    /// Set read timeout in milliseconds (i.e. the amount of time that the `read_packet()` method
    /// can block).
    pub fn read_timeout(self, read_timeout: u64) -> Self {
        unsafe {
            pcap_wrapper__set_read_timeout(self.inner.ptr, read_timeout);
        }

        self
    }

    /// Set packet filter.
    pub fn filter(self, filter: Option<&str>) -> Self {
        let filter = filter.map(|s| CString::new(s).unwrap());

        let filter_ptr = filter.as_ref().map(|f| f.as_ptr()).unwrap_or(ptr::null());

        let ret = unsafe { pcap_wrapper__set_filter(self.inner.ptr, filter_ptr) };

        if ret != 0 {
            panic!("Unable to allocate a packet filter");
        }

        self
    }

    /// Activate the capture.
    pub fn activate(self) -> Result<Capture> {
        let _tc = TC.lock().unwrap();

        let ret = unsafe { pcap_wrapper__open(self.inner.ptr) };

        if ret == 0 {
            Ok(self.inner)
        } else {
            Err(self.inner.get_last_error())
        }
    }
}

/// Capture handle that can be used for reading and writing raw packets.
struct Capture {
    ptr: *mut c_void,
}

impl Capture {
    /// Get a builder.
    pub fn builder(device: &str) -> CaptureBuilder {
        CaptureBuilder::new(device)
    }

    /// Read one packet. The packet will be appended to a given buffer. True will be returned if a
    /// packet was read, false will be returned in case of timeout.
    pub fn read_packet(&mut self, buffer: &mut BytesMut) -> Result<bool> {
        let buffer_ptr: *mut BytesMut = buffer;

        let ret =
            unsafe { pcap_wrapper__read_packet(self.ptr, read_packet_callback, buffer_ptr as _) };

        if ret >= 0 {
            Ok(ret > 0)
        } else {
            Err(self.get_last_error())
        }
    }

    /// Write a given packet to the underlying network device.
    pub fn write_packet(&mut self, packet: &[u8]) -> Result<()> {
        let data = packet.as_ptr();
        let size = packet.len();

        let ret = unsafe { pcap_wrapper__write_packet(self.ptr, data, size as _) };

        if ret == 0 {
            Ok(())
        } else {
            Err(self.get_last_error())
        }
    }

    /// Get last error.
    fn get_last_error(&self) -> PcapError {
        unsafe {
            let ptr = pcap_wrapper__get_last_error(self.ptr);
            let msg = CStr::from_ptr(ptr);

            PcapError::new(msg.to_string_lossy())
        }
    }
}

impl Drop for Capture {
    fn drop(&mut self) {
        unsafe { pcap_wrapper__free(self.ptr) }
    }
}

/// Helper callback for getting packet data.
unsafe extern "C" fn read_packet_callback(
    opaque: *mut c_void,
    data: *const u8,
    data_length: size_t,
    _: size_t,
) {
    let buffer_ptr = opaque as *mut BytesMut;

    let data = slice::from_raw_parts(data, data_length as _);

    (*buffer_ptr).extend_from_slice(data);
}

/// PCAP packet scanner (implementation of a send-receive service).
pub struct Scanner {
    device: String,
}

impl Scanner {
    /// Create a new PCAP scanner for a given device.
    pub fn new(device: &str) -> Self {
        Self {
            device: device.to_string(),
        }
    }

    /// Send all packets from a given iterator and receive all packets
    /// according to a given filter.
    pub fn sr<F>(
        &mut self,
        filter: &str,
        packet_generator: F,
        read_timeout: Duration,
        stop_after: Option<Duration>,
    ) -> Result<Vec<EtherPacket>>
    where
        F: FnMut() -> Option<Bytes>,
    {
        let device = self.device.clone();
        let filter = filter.to_string();

        // NOTE: we need to synchronize the sender and the receiver parts to avoid sending out
        // packets before the receiver is initialized (otherwise we could miss some response
        // packets)
        let (tx, rx) = mpsc::channel();

        let t = thread::spawn(move || {
            packet_listener(&device, &filter, read_timeout, stop_after, Some(tx))
        });

        rx.recv().unwrap_or_default();

        send_packets(&self.device, packet_generator)?;

        t.join().unwrap()
    }
}

/// Listen for incoming packets matching a given filter until no packets are read for a given
/// soft timeout or until a given hard timeout expires.
///
/// # Arguments
/// * `device` - device name
/// * `filter` - packet filter expression
/// * `packet_timeout` - a soft timeout that will stop the listener if no packet is received for
///   given duration
/// * `total_timeout` - a hard timeout that will stop the listener when it expires
/// * `init_event_tx` - an optional sender that will be used to notify another thread that the
///   listener has been initialized and is receiving packets
fn packet_listener(
    device: &str,
    filter: &str,
    packet_timeout: Duration,
    total_timeout: Option<Duration>,
    init_event_tx: Option<Sender<()>>,
) -> Result<Vec<EtherPacket>> {
    let mut cap = Capture::builder(device)
        .max_packet_length(65_536)
        .read_timeout(100)
        .filter(Some(filter))
        .activate()?;

    if let Some(tx) = init_event_tx {
        tx.send(()).unwrap_or_default();
    }

    let start_time = Instant::now();

    let mut last_packet_time = Instant::now();
    let mut buffer = BytesMut::new();
    let mut res = Vec::new();

    loop {
        if cap.read_packet(&mut buffer)? {
            last_packet_time = Instant::now();

            if let Ok(packet) = EtherPacket::parse(buffer.as_ref()) {
                res.push(packet);
            }

            buffer.clear();
        }

        if let Some(timeout) = total_timeout {
            if start_time.elapsed() > timeout {
                break;
            }
        }

        if last_packet_time.elapsed() > packet_timeout {
            break;
        }
    }

    Ok(res)
}

/// Send all packets using a given network device.
fn send_packets<F>(device: &str, mut packet_generator: F) -> Result<()>
where
    F: FnMut() -> Option<Bytes>,
{
    let mut cap = Capture::builder(device).activate()?;

    while let Some(pkt) = packet_generator() {
        cap.write_packet(pkt.as_ref())?;
    }

    Ok(())
}
<commit_msg>Fix build without libc<commit_after>// Copyright 2015 click2stream, Inc.
//
// 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.

//! PCAP network scanner definitions.

use std::fmt;
use std::ptr;
use std::slice;
use std::sync::mpsc;
use std::thread;

use std::error::Error;
use std::ffi::{CStr, CString};
use std::fmt::{Display, Formatter};
use std::os::raw::{c_char, c_int, c_void};
use std::sync::mpsc::Sender;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use bytes::{Bytes, BytesMut};

use crate::net::raw::ether::packet::EtherPacket;

type PacketCallback = unsafe extern "C" fn(
    opaque: *mut c_void,
    data: *const u8,
    data_length: usize,
    packet_length: usize,
);

extern "C" {
    fn pcap_wrapper__new(device: *const c_char) -> *mut c_void;
    fn pcap_wrapper__free(wrapper: *mut c_void);

    fn pcap_wrapper__get_last_error(wrapper: *const c_void) -> *const c_char;

    fn pcap_wrapper__set_filter(wrapper: *mut c_void, filter: *const c_char) -> c_int;
    fn pcap_wrapper__set_max_packet_length(wrapper: *mut c_void, max_packet_length: usize);
    fn pcap_wrapper__set_read_timeout(wrapper: *mut c_void, read_timeout: u64);

    fn pcap_wrapper__open(wrapper: *mut c_void) -> c_int;

    fn pcap_wrapper__read_packet(
        wrapper: *mut c_void,
        callback: PacketCallback,
        opaque: *mut c_void,
    ) -> c_int;
    fn pcap_wrapper__write_packet(wrapper: *mut c_void, data: *const u8, size: usize) -> c_int;
}

lazy_static! {
    /// PCAP context for synchronizing thread unsafe calls.
    static ref TC: Mutex<()> = Mutex::new(());
}

/// PCAP error.
#[derive(Debug, Clone)]
pub struct PcapError {
    msg: String,
}

impl PcapError {
    /// Create a new error with a given error message.
    pub fn new<T>(msg: T) -> Self
    where
        T: ToString,
    {
        Self {
            msg: msg.to_string(),
        }
    }
}

impl Display for PcapError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        f.write_str(&self.msg)
    }
}

impl Error for PcapError {}

/// Common result type for this module.
pub type Result<T> = std::result::Result<T, PcapError>;

/// Capture builder.
struct CaptureBuilder {
    inner: Capture,
}

impl CaptureBuilder {
    /// Create a new capture builder for a given device.
    fn new(device: &str) -> Self {
        let device = CString::new(device).unwrap();

        let ptr = unsafe { pcap_wrapper__new(device.as_ptr()) };

        if ptr.is_null() {
            panic!("Unable to allocate PCAP capture device");
        }

        let capture = Capture { ptr };

        Self { inner: capture }
    }

    /// Set maximum packet length (longer packets will be truncated).
    pub fn max_packet_length(self, max_packet_length: usize) -> Self {
        unsafe {
            pcap_wrapper__set_max_packet_length(self.inner.ptr, max_packet_length as _);
        }

        self
    }

    /// Set read timeout in milliseconds (i.e. the amount of time that the `read_packet()` method
    /// can block).
    pub fn read_timeout(self, read_timeout: u64) -> Self {
        unsafe {
            pcap_wrapper__set_read_timeout(self.inner.ptr, read_timeout);
        }

        self
    }

    /// Set packet filter.
    pub fn filter(self, filter: Option<&str>) -> Self {
        let filter = filter.map(|s| CString::new(s).unwrap());

        let filter_ptr = filter.as_ref().map(|f| f.as_ptr()).unwrap_or(ptr::null());

        let ret = unsafe { pcap_wrapper__set_filter(self.inner.ptr, filter_ptr) };

        if ret != 0 {
            panic!("Unable to allocate a packet filter");
        }

        self
    }

    /// Activate the capture.
    pub fn activate(self) -> Result<Capture> {
        let _tc = TC.lock().unwrap();

        let ret = unsafe { pcap_wrapper__open(self.inner.ptr) };

        if ret == 0 {
            Ok(self.inner)
        } else {
            Err(self.inner.get_last_error())
        }
    }
}

/// Capture handle that can be used for reading and writing raw packets.
struct Capture {
    ptr: *mut c_void,
}

impl Capture {
    /// Get a builder.
    pub fn builder(device: &str) -> CaptureBuilder {
        CaptureBuilder::new(device)
    }

    /// Read one packet. The packet will be appended to a given buffer. True will be returned if a
    /// packet was read, false will be returned in case of timeout.
    pub fn read_packet(&mut self, buffer: &mut BytesMut) -> Result<bool> {
        let buffer_ptr: *mut BytesMut = buffer;

        let ret =
            unsafe { pcap_wrapper__read_packet(self.ptr, read_packet_callback, buffer_ptr as _) };

        if ret >= 0 {
            Ok(ret > 0)
        } else {
            Err(self.get_last_error())
        }
    }

    /// Write a given packet to the underlying network device.
    pub fn write_packet(&mut self, packet: &[u8]) -> Result<()> {
        let data = packet.as_ptr();
        let size = packet.len();

        let ret = unsafe { pcap_wrapper__write_packet(self.ptr, data, size as _) };

        if ret == 0 {
            Ok(())
        } else {
            Err(self.get_last_error())
        }
    }

    /// Get last error.
    fn get_last_error(&self) -> PcapError {
        unsafe {
            let ptr = pcap_wrapper__get_last_error(self.ptr);
            let msg = CStr::from_ptr(ptr);

            PcapError::new(msg.to_string_lossy())
        }
    }
}

impl Drop for Capture {
    fn drop(&mut self) {
        unsafe { pcap_wrapper__free(self.ptr) }
    }
}

/// Helper callback for getting packet data.
unsafe extern "C" fn read_packet_callback(
    opaque: *mut c_void,
    data: *const u8,
    data_length: usize,
    _: usize,
) {
    let buffer_ptr = opaque as *mut BytesMut;

    let data = slice::from_raw_parts(data, data_length as _);

    (*buffer_ptr).extend_from_slice(data);
}

/// PCAP packet scanner (implementation of a send-receive service).
pub struct Scanner {
    device: String,
}

impl Scanner {
    /// Create a new PCAP scanner for a given device.
    pub fn new(device: &str) -> Self {
        Self {
            device: device.to_string(),
        }
    }

    /// Send all packets from a given iterator and receive all packets
    /// according to a given filter.
    pub fn sr<F>(
        &mut self,
        filter: &str,
        packet_generator: F,
        read_timeout: Duration,
        stop_after: Option<Duration>,
    ) -> Result<Vec<EtherPacket>>
    where
        F: FnMut() -> Option<Bytes>,
    {
        let device = self.device.clone();
        let filter = filter.to_string();

        // NOTE: we need to synchronize the sender and the receiver parts to avoid sending out
        // packets before the receiver is initialized (otherwise we could miss some response
        // packets)
        let (tx, rx) = mpsc::channel();

        let t = thread::spawn(move || {
            packet_listener(&device, &filter, read_timeout, stop_after, Some(tx))
        });

        rx.recv().unwrap_or_default();

        send_packets(&self.device, packet_generator)?;

        t.join().unwrap()
    }
}

/// Listen for incoming packets matching a given filter until no packets are read for a given
/// soft timeout or until a given hard timeout expires.
///
/// # Arguments
/// * `device` - device name
/// * `filter` - packet filter expression
/// * `packet_timeout` - a soft timeout that will stop the listener if no packet is received for
///   given duration
/// * `total_timeout` - a hard timeout that will stop the listener when it expires
/// * `init_event_tx` - an optional sender that will be used to notify another thread that the
///   listener has been initialized and is receiving packets
fn packet_listener(
    device: &str,
    filter: &str,
    packet_timeout: Duration,
    total_timeout: Option<Duration>,
    init_event_tx: Option<Sender<()>>,
) -> Result<Vec<EtherPacket>> {
    let mut cap = Capture::builder(device)
        .max_packet_length(65_536)
        .read_timeout(100)
        .filter(Some(filter))
        .activate()?;

    if let Some(tx) = init_event_tx {
        tx.send(()).unwrap_or_default();
    }

    let start_time = Instant::now();

    let mut last_packet_time = Instant::now();
    let mut buffer = BytesMut::new();
    let mut res = Vec::new();

    loop {
        if cap.read_packet(&mut buffer)? {
            last_packet_time = Instant::now();

            if let Ok(packet) = EtherPacket::parse(buffer.as_ref()) {
                res.push(packet);
            }

            buffer.clear();
        }

        if let Some(timeout) = total_timeout {
            if start_time.elapsed() > timeout {
                break;
            }
        }

        if last_packet_time.elapsed() > packet_timeout {
            break;
        }
    }

    Ok(res)
}

/// Send all packets using a given network device.
fn send_packets<F>(device: &str, mut packet_generator: F) -> Result<()>
where
    F: FnMut() -> Option<Bytes>,
{
    let mut cap = Capture::builder(device).activate()?;

    while let Some(pkt) = packet_generator() {
        cap.write_packet(pkt.as_ref())?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>use prolog::ast::*;

use std::collections::HashMap;

pub type PredicateKey = (Atom, usize); // name, arity, type.

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub enum PredicateKeyType {
    BuiltIn,
    User
}

pub type OpDirKey = (Atom, Fixity);
// name and fixity -> operator type and precedence.
pub type OpDir = HashMap<OpDirKey, (Specifier, usize)>;

pub type CodeDir = HashMap<PredicateKey, (PredicateKeyType, usize)>;

fn get_builtins() -> Code {
    vec![internal_call_n!(), // callN/N, 0.
         is_atomic!(temp_v!(1)), // atomic/1, 1.
         proceed!(),
         is_var!(temp_v!(1)),    // var/1, 3.
         proceed!(),
         allocate!(4), // catch/3, 5.
         fact![get_var_in_fact!(perm_v!(2), 1),
               get_var_in_fact!(perm_v!(3), 2),
               get_var_in_fact!(perm_v!(1), 3)],
         query![put_var!(perm_v!(4), 1)],
         get_current_block!(),
         query![put_value!(perm_v!(2), 1),
                put_value!(perm_v!(3), 2),
                put_value!(perm_v!(1), 3),
                put_unsafe_value!(4, 4)],
         deallocate!(),
         goto!(12, 4), // goto catch/4.
         try_me_else!(10), // catch/4, 12.
         allocate!(3),
         fact![get_var_in_fact!(perm_v!(3), 1),
               get_var_in_fact!(perm_v!(2), 4)],
         query![put_var!(perm_v!(1), 1)],
         install_new_block!(),
         query![put_value!(perm_v!(3), 1)],
         call_n!(1),
         query![put_value!(perm_v!(2), 1),
                put_unsafe_value!(1, 2)],
         deallocate!(),
         goto!(44, 2), //21: goto end_block/2.
         trust_me!(),
         allocate!(3),
         fact![get_var_in_fact!(perm_v!(2), 2),
               get_var_in_fact!(perm_v!(1), 3)],
         query![get_var_in_query!(temp_v!(2), 1),
                put_value!(temp_v!(4), 1)],
         reset_block!(),
         query![put_var!(perm_v!(3), 1)],
         get_ball!(),
         query![put_unsafe_value!(3, 1),
                put_value!(perm_v!(2), 2),
                put_value!(perm_v!(1), 3)],
         deallocate!(),
         goto!(32, 2), // goto handle_ball/2.
         try_me_else!(10), // handle_ball/2, 32.
         allocate!(2),
         get_level!(),
         fact![get_var_in_fact!(perm_v!(2), 3)],
         unify!(),
         cut!(),
         erase_ball!(),
         query![put_value!(perm_v!(2), 1)],
         deallocate!(),
         execute_n!(1),
         trust_me!(),
         unwind_stack!(),
         try_me_else!(9), // end_block/2, 44.
         allocate!(1),
         fact![get_var_in_fact!(perm_v!(1), 1)],
         query![put_value!(temp_v!(2), 1)],
         clean_up_block!(),
         query![put_value!(perm_v!(1), 1)],
         deallocate!(),
         reset_block!(),
         proceed!(),
         trust_me!(), // 53.
         allocate!(0),
         query![get_var_in_query!(temp_v!(3), 1),
                put_value!(temp_v!(2), 1)],
         reset_block!(),
         deallocate!(),
         goto!(61, 0),
         set_ball!(), // throw/1, 59.
         unwind_stack!(),
         fail!(), // false/0, 61.
         try_me_else!(7), // not/1, 62.
         allocate!(1),
         get_level!(),
         call_n!(1),
         cut!(),
         deallocate!(),
         goto!(61, 0),
         trust_me!(),
         proceed!(),
         duplicate_term!(), // duplicate_term/2, 71.
         proceed!(),
         fact![get_value!(temp_v!(1), 2)], // =/2, 73.
         proceed!(),
         proceed!(), // true/0, 75.
         try_me_else!(2), // ';'/2, 76.
         execute_n!(1),
         trust_me!(),
         query![put_value!(temp_v!(2), 1)],
         execute_n!(1),
         get_cp!(temp_v!(3)), // ','/2, 81
         goto!(83, 3),
         try_me_else!(18), // ','/3, 83.
         switch_on_term!(4, 1, 0, 0),
         indexed_try!(4),
         retry!(7),
         trust!(10),
         try_me_else!(4),
         fact![get_constant!(Constant::from("!"), temp_v!(1)),
               get_structure!(String::from(","), 2, temp_v!(2)),
               unify_variable!(temp_v!(1)),
               unify_variable!(temp_v!(2))],
         set_neck_cut!(temp_v!(3)),
         goto!(83, 3),
         retry_me_else!(4),
         fact![get_constant!(Constant::from("!"), temp_v!(1)),
               get_constant!(Constant::from("!"), temp_v!(2))],
         set_neck_cut!(temp_v!(3)),
         proceed!(),
         trust_me!(),
         fact![get_constant!(Constant::from("!"), temp_v!(1))],
         set_neck_cut!(temp_v!(3)),
         query![put_value!(temp_v!(2), 1)],
         execute_n!(1),
         retry_me_else!(8),
         allocate!(2),
         fact![get_structure!(String::from(","), 2, temp_v!(2)),
               unify_variable!(perm_v!(2)),
               unify_variable!(perm_v!(1))],
         neck_cut!(),
         call_n!(1),
         query![put_unsafe_value!(2, 1),
                put_unsafe_value!(1, 2)],
         deallocate!(),
         goto!(83, 3),
         retry_me_else!(9),
         allocate!(1),
         get_level!(),
         fact![get_constant!(Constant::from("!"), temp_v!(2))],
         neck_cut!(),
         call_n!(1),
         set_non_neck_cut!(temp_v!(3)),
         deallocate!(),
         proceed!(),
         trust_me!(),
         allocate!(1),
         fact![get_var_in_fact!(perm_v!(1), 2)],
         call_n!(1),
         query![put_value!(perm_v!(1), 1)],
         deallocate!(),
         execute_n!(1),
         allocate!(2), // (->)/2, 124.
         get_level!(),
         fact![get_var_in_fact!(perm_v!(2), 2)],
         call_n!(1),
         cut!(),
         query![put_value!(perm_v!(2), 1)],
         deallocate!(),
         execute_n!(1)
    ]
}

pub fn build_code_dir() -> (Code, CodeDir, OpDir)
{
    let mut code_dir = HashMap::new();
    let mut op_dir   = HashMap::new();

    let builtin_code = get_builtins();

    op_dir.insert((String::from(":-"), Fixity::In),   (XFX, 1200));
    op_dir.insert((String::from(":-"), Fixity::Pre),  (FX, 1200));
    op_dir.insert((String::from("?-"), Fixity::Pre),  (FX, 1200));

    // control operators.
    op_dir.insert((String::from("\\+"), Fixity::Pre), (FY, 900));
    op_dir.insert((String::from("="), Fixity::In), (XFX, 700));

    // arithmetic operators.
    op_dir.insert((String::from("is"), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("+"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("-"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("/\\"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("\\/"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("xor"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("//"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("/"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("div"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("*"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("-"), Fixity::Pre), (FY, 200));
    op_dir.insert((String::from("rdiv"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("<<"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from(">>"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("mod"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("rem"), Fixity::In), (YFX, 400));

    // arithmetic comparison operators.
    op_dir.insert((String::from(">"), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("<"), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("=\\="), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("=:="), Fixity::In), (XFX, 700));
    op_dir.insert((String::from(">="), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("=<"), Fixity::In), (XFX, 700));

    // control operators.
    op_dir.insert((String::from(";"), Fixity::In), (XFY, 1100));
    op_dir.insert((String::from("->"), Fixity::In), (XFY, 1050));

    // there are 63 registers in the VM, so call/N is defined for all 0 <= N <= 62
    // (an extra register is needed for the predicate name)
    for arity in 0 .. 63 {
        code_dir.insert((String::from("call"), arity), (PredicateKeyType::BuiltIn, 0));
    }

    code_dir.insert((String::from("atomic"), 1), (PredicateKeyType::BuiltIn, 1));
    code_dir.insert((String::from("var"), 1), (PredicateKeyType::BuiltIn, 3));
    code_dir.insert((String::from("false"), 0), (PredicateKeyType::BuiltIn, 61));
    code_dir.insert((String::from("\\+"), 1), (PredicateKeyType::BuiltIn, 62));
    code_dir.insert((String::from("duplicate_term"), 2), (PredicateKeyType::BuiltIn, 71));
    code_dir.insert((String::from("catch"), 3), (PredicateKeyType::BuiltIn, 5));
    code_dir.insert((String::from("throw"), 1), (PredicateKeyType::BuiltIn, 59));
    code_dir.insert((String::from("="), 2), (PredicateKeyType::BuiltIn, 73));
    code_dir.insert((String::from("true"), 0), (PredicateKeyType::BuiltIn, 75));

    code_dir.insert((String::from(";"), 2), (PredicateKeyType::BuiltIn, 76));
    code_dir.insert((String::from(","), 2), (PredicateKeyType::BuiltIn, 81));
    code_dir.insert((String::from("->"), 2), (PredicateKeyType::BuiltIn, 124));

    (builtin_code, code_dir, op_dir)
}
<commit_msg>correction to -> offset<commit_after>use prolog::ast::*;

use std::collections::HashMap;

pub type PredicateKey = (Atom, usize); // name, arity, type.

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub enum PredicateKeyType {
    BuiltIn,
    User
}

pub type OpDirKey = (Atom, Fixity);
// name and fixity -> operator type and precedence.
pub type OpDir = HashMap<OpDirKey, (Specifier, usize)>;

pub type CodeDir = HashMap<PredicateKey, (PredicateKeyType, usize)>;

fn get_builtins() -> Code {
    vec![internal_call_n!(), // callN/N, 0.
         is_atomic!(temp_v!(1)), // atomic/1, 1.
         proceed!(),
         is_var!(temp_v!(1)),    // var/1, 3.
         proceed!(),
         allocate!(4), // catch/3, 5.
         fact![get_var_in_fact!(perm_v!(2), 1),
               get_var_in_fact!(perm_v!(3), 2),
               get_var_in_fact!(perm_v!(1), 3)],
         query![put_var!(perm_v!(4), 1)],
         get_current_block!(),
         query![put_value!(perm_v!(2), 1),
                put_value!(perm_v!(3), 2),
                put_value!(perm_v!(1), 3),
                put_unsafe_value!(4, 4)],
         deallocate!(),
         goto!(12, 4), // goto catch/4.
         try_me_else!(10), // catch/4, 12.
         allocate!(3),
         fact![get_var_in_fact!(perm_v!(3), 1),
               get_var_in_fact!(perm_v!(2), 4)],
         query![put_var!(perm_v!(1), 1)],
         install_new_block!(),
         query![put_value!(perm_v!(3), 1)],
         call_n!(1),
         query![put_value!(perm_v!(2), 1),
                put_unsafe_value!(1, 2)],
         deallocate!(),
         goto!(44, 2), //21: goto end_block/2.
         trust_me!(),
         allocate!(3),
         fact![get_var_in_fact!(perm_v!(2), 2),
               get_var_in_fact!(perm_v!(1), 3)],
         query![get_var_in_query!(temp_v!(2), 1),
                put_value!(temp_v!(4), 1)],
         reset_block!(),
         query![put_var!(perm_v!(3), 1)],
         get_ball!(),
         query![put_unsafe_value!(3, 1),
                put_value!(perm_v!(2), 2),
                put_value!(perm_v!(1), 3)],
         deallocate!(),
         goto!(32, 2), // goto handle_ball/2.
         try_me_else!(10), // handle_ball/2, 32.
         allocate!(2),
         get_level!(),
         fact![get_var_in_fact!(perm_v!(2), 3)],
         unify!(),
         cut!(),
         erase_ball!(),
         query![put_value!(perm_v!(2), 1)],
         deallocate!(),
         execute_n!(1),
         trust_me!(),
         unwind_stack!(),
         try_me_else!(9), // end_block/2, 44.
         allocate!(1),
         fact![get_var_in_fact!(perm_v!(1), 1)],
         query![put_value!(temp_v!(2), 1)],
         clean_up_block!(),
         query![put_value!(perm_v!(1), 1)],
         deallocate!(),
         reset_block!(),
         proceed!(),
         trust_me!(), // 53.
         allocate!(0),
         query![get_var_in_query!(temp_v!(3), 1),
                put_value!(temp_v!(2), 1)],
         reset_block!(),
         deallocate!(),
         goto!(61, 0),
         set_ball!(), // throw/1, 59.
         unwind_stack!(),
         fail!(), // false/0, 61.
         try_me_else!(7), // not/1, 62.
         allocate!(1),
         get_level!(),
         call_n!(1),
         cut!(),
         deallocate!(),
         goto!(61, 0),
         trust_me!(),
         proceed!(),
         duplicate_term!(), // duplicate_term/2, 71.
         proceed!(),
         fact![get_value!(temp_v!(1), 2)], // =/2, 73.
         proceed!(),
         proceed!(), // true/0, 75.
         try_me_else!(2), // ';'/2, 76.
         execute_n!(1),
         trust_me!(),
         query![put_value!(temp_v!(2), 1)],
         execute_n!(1),
         get_cp!(temp_v!(3)), // ','/2, 81
         goto!(83, 3),
         try_me_else!(18), // ','/3, 83.
         switch_on_term!(4, 1, 0, 0),
         indexed_try!(4),
         retry!(7),
         trust!(10),
         try_me_else!(4),
         fact![get_constant!(Constant::from("!"), temp_v!(1)),
               get_structure!(String::from(","), 2, temp_v!(2)),
               unify_variable!(temp_v!(1)),
               unify_variable!(temp_v!(2))],
         set_neck_cut!(temp_v!(3)),
         goto!(83, 3),
         retry_me_else!(4),
         fact![get_constant!(Constant::from("!"), temp_v!(1)),
               get_constant!(Constant::from("!"), temp_v!(2))],
         set_neck_cut!(temp_v!(3)),
         proceed!(),
         trust_me!(),
         fact![get_constant!(Constant::from("!"), temp_v!(1))],
         set_neck_cut!(temp_v!(3)),
         query![put_value!(temp_v!(2), 1)],
         execute_n!(1),
         retry_me_else!(8),
         allocate!(2),
         fact![get_structure!(String::from(","), 2, temp_v!(2)),
               unify_variable!(perm_v!(2)),
               unify_variable!(perm_v!(1))],
         neck_cut!(),
         call_n!(1),
         query![put_unsafe_value!(2, 1),
                put_unsafe_value!(1, 2)],
         deallocate!(),
         goto!(83, 3),
         retry_me_else!(9),
         allocate!(1),
         get_level!(),
         fact![get_constant!(Constant::from("!"), temp_v!(2))],
         neck_cut!(),
         call_n!(1),
         set_non_neck_cut!(temp_v!(3)),
         deallocate!(),
         proceed!(),
         trust_me!(),
         allocate!(1),
         fact![get_var_in_fact!(perm_v!(1), 2)],
         call_n!(1),
         query![put_value!(perm_v!(1), 1)],
         deallocate!(),
         execute_n!(1),
         allocate!(2), // (->)/2, 125.
         get_level!(),
         fact![get_var_in_fact!(perm_v!(2), 2)],
         call_n!(1),
         cut!(),
         query![put_value!(perm_v!(2), 1)],
         deallocate!(),
         execute_n!(1)
    ]
}

pub fn build_code_dir() -> (Code, CodeDir, OpDir)
{
    let mut code_dir = HashMap::new();
    let mut op_dir   = HashMap::new();

    let builtin_code = get_builtins();

    op_dir.insert((String::from(":-"), Fixity::In),   (XFX, 1200));
    op_dir.insert((String::from(":-"), Fixity::Pre),  (FX, 1200));
    op_dir.insert((String::from("?-"), Fixity::Pre),  (FX, 1200));

    // control operators.
    op_dir.insert((String::from("\\+"), Fixity::Pre), (FY, 900));
    op_dir.insert((String::from("="), Fixity::In), (XFX, 700));

    // arithmetic operators.
    op_dir.insert((String::from("is"), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("+"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("-"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("/\\"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("\\/"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("xor"), Fixity::In), (YFX, 500));
    op_dir.insert((String::from("//"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("/"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("div"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("*"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("-"), Fixity::Pre), (FY, 200));
    op_dir.insert((String::from("rdiv"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("<<"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from(">>"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("mod"), Fixity::In), (YFX, 400));
    op_dir.insert((String::from("rem"), Fixity::In), (YFX, 400));

    // arithmetic comparison operators.
    op_dir.insert((String::from(">"), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("<"), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("=\\="), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("=:="), Fixity::In), (XFX, 700));
    op_dir.insert((String::from(">="), Fixity::In), (XFX, 700));
    op_dir.insert((String::from("=<"), Fixity::In), (XFX, 700));

    // control operators.
    op_dir.insert((String::from(";"), Fixity::In), (XFY, 1100));
    op_dir.insert((String::from("->"), Fixity::In), (XFY, 1050));

    // there are 63 registers in the VM, so call/N is defined for all 0 <= N <= 62
    // (an extra register is needed for the predicate name)
    for arity in 0 .. 63 {
        code_dir.insert((String::from("call"), arity), (PredicateKeyType::BuiltIn, 0));
    }

    code_dir.insert((String::from("atomic"), 1), (PredicateKeyType::BuiltIn, 1));
    code_dir.insert((String::from("var"), 1), (PredicateKeyType::BuiltIn, 3));
    code_dir.insert((String::from("false"), 0), (PredicateKeyType::BuiltIn, 61));
    code_dir.insert((String::from("\\+"), 1), (PredicateKeyType::BuiltIn, 62));
    code_dir.insert((String::from("duplicate_term"), 2), (PredicateKeyType::BuiltIn, 71));
    code_dir.insert((String::from("catch"), 3), (PredicateKeyType::BuiltIn, 5));
    code_dir.insert((String::from("throw"), 1), (PredicateKeyType::BuiltIn, 59));
    code_dir.insert((String::from("="), 2), (PredicateKeyType::BuiltIn, 73));
    code_dir.insert((String::from("true"), 0), (PredicateKeyType::BuiltIn, 75));

    code_dir.insert((String::from(";"), 2), (PredicateKeyType::BuiltIn, 76));
    code_dir.insert((String::from(","), 2), (PredicateKeyType::BuiltIn, 81));
    code_dir.insert((String::from("->"), 2), (PredicateKeyType::BuiltIn, 125));

    (builtin_code, code_dir, op_dir)
}
<|endoftext|>


<commit_before>//! This module defines graphics pipelines and types bound to those pipelines.
//!
//! Namely, this includes vertex and uniform buffer types.

use gfx;
use std::hash::{Hash, Hasher};

use math::{FMatrix4, FPoint2, FPoint3, Matrix4Ext};
use super::Color;

pub use self::pipeline::*;

gfx_pipeline!{
    pipeline {
        buffer: gfx::VertexBuffer<Vertex> = (),
        transform: gfx::ConstantBuffer<Transform> = "transform",
        camera: gfx::Global<[[f32; 4]; 4]> = "u_camera",
        model: gfx::Global<[[f32; 4]; 4]> = "u_model",
        sampler: gfx::TextureSampler<[f32; 4]> = "t_texture",
        color: gfx::RenderTarget<gfx::format::Rgba8> = "f_target0",
        depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
    }
}

gfx_constant_struct!{
    Transform {
        camera: [[f32; 4]; 4] = "u_camera",
        model: [[f32; 4]; 4] = "u_model",
    }
}

impl Transform {
    pub fn new(camera: &FMatrix4, model: &FMatrix4) -> Self {
        Transform {
            camera: camera.to_array(),
            model: model.to_array(),
        }
    }
}

impl Default for Transform {
    fn default() -> Self {
        let identity = FMatrix4::identity();
        Transform::new(&identity, &identity)
    }
}

gfx_vertex_struct!{
    Vertex {
        position: [f32; 3] = "a_position",
        uv: [f32; 2] = "a_uv",
        color: [f32; 4] = "a_color",
    }
}

impl Vertex {
    pub fn new(position: &FPoint3, uv: &FPoint2, color: &Color) -> Self {
        Vertex {
            position: [position.x, position.y, position.z],
            uv: [uv.x, uv.y],
            color: [color.x, color.y, color.z, color.w],
        }
    }
}

impl Default for Vertex {
    fn default() -> Self {
        Vertex::new(
            &FPoint3::origin(),
            &FPoint2::origin(),
            &Color::new(0.0, 0.0, 0.0, 1.0),
        )
    }
}

// TODO: This violates the contract of the `Eq` trait, because `Vertex`
//       includes floating point data. For example, NaN is problematic. See the
//       implementation of `Hash`.
impl Eq for Vertex {}

impl Hash for Vertex {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        use std::mem;

        // TODO: This is awful. See the ordered-float crate for a far better
        //       approach to hashing floats.
        unsafe {
            mem::transmute::<[f32; 3], [u32; 3]>(self.position).hash(state);
            mem::transmute::<[f32; 2], [u32; 2]>(self.uv).hash(state);
            mem::transmute::<[f32; 4], [u32; 4]>(self.color).hash(state);
        }
    }
}

impl PartialEq for Vertex {
    fn eq(&self, other: &Self) -> bool {
        (self.position == other.position) && (self.uv == other.uv) && (self.color == other.color)
    }
}
<commit_msg>Use `OrderedFloat` to hash `Vertex`s.<commit_after>//! This module defines graphics pipelines and types bound to those pipelines.
//!
//! Namely, this includes vertex and uniform buffer types.

use gfx;
use plexus::r32;
use std::hash::{Hash, Hasher};

use math::{FMatrix4, FPoint2, FPoint3, Matrix4Ext};
use super::Color;

pub use self::pipeline::*;

gfx_pipeline!{
    pipeline {
        buffer: gfx::VertexBuffer<Vertex> = (),
        transform: gfx::ConstantBuffer<Transform> = "transform",
        camera: gfx::Global<[[f32; 4]; 4]> = "u_camera",
        model: gfx::Global<[[f32; 4]; 4]> = "u_model",
        sampler: gfx::TextureSampler<[f32; 4]> = "t_texture",
        color: gfx::RenderTarget<gfx::format::Rgba8> = "f_target0",
        depth: gfx::DepthTarget<gfx::format::DepthStencil> = gfx::preset::depth::LESS_EQUAL_WRITE,
    }
}

gfx_constant_struct!{
    Transform {
        camera: [[f32; 4]; 4] = "u_camera",
        model: [[f32; 4]; 4] = "u_model",
    }
}

impl Transform {
    pub fn new(camera: &FMatrix4, model: &FMatrix4) -> Self {
        Transform {
            camera: camera.to_array(),
            model: model.to_array(),
        }
    }
}

impl Default for Transform {
    fn default() -> Self {
        let identity = FMatrix4::identity();
        Transform::new(&identity, &identity)
    }
}

gfx_vertex_struct!{
    Vertex {
        position: [f32; 3] = "a_position",
        uv: [f32; 2] = "a_uv",
        color: [f32; 4] = "a_color",
    }
}

impl Vertex {
    pub fn new(position: &FPoint3, uv: &FPoint2, color: &Color) -> Self {
        Vertex {
            position: [position.x, position.y, position.z],
            uv: [uv.x, uv.y],
            color: [color.x, color.y, color.z, color.w],
        }
    }

    fn ordered(&self) -> OrderedVertex {
        OrderedVertex {
            position: [
                r32::from(self.position[0]),
                r32::from(self.position[1]),
                r32::from(self.position[2]),
            ],
            uv: [r32::from(self.uv[0]), r32::from(self.uv[1])],
            color: [
                r32::from(self.color[0]),
                r32::from(self.color[1]),
                r32::from(self.color[2]),
                r32::from(self.color[3]),
            ],
        }
    }
}

impl Default for Vertex {
    fn default() -> Self {
        Vertex::new(
            &FPoint3::origin(),
            &FPoint2::origin(),
            &Color::new(0.0, 0.0, 0.0, 1.0),
        )
    }
}

impl Eq for Vertex {}

impl Hash for Vertex {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        self.ordered().hash(state);
    }
}

impl PartialEq for Vertex {
    fn eq(&self, other: &Self) -> bool {
        self.ordered().eq(&other.ordered())
    }
}

#[derive(Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
struct OrderedVertex {
    position: [r32; 3],
    uv: [r32; 2],
    color: [r32; 4],
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use ipc_channel::ipc::IpcError;
use serde::{Serialize, Deserialize};
use tokio::sync::{mpsc, RwLock, Mutex};
use thiserror::Error;

use crate::ipc_protocol::{ClientConnection, ConnectionError, ClientRequest, ServerResponse};
use crate::renderer_server::RendererServer;

/// Signals that the IPC connection has been disconnected and therefore the window was probably
/// closed
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
#[error("Cannot continue to run turtle commands after window is closed. This panic stops the thread, but is not necessarily an error.")]
struct Disconnected;

/// A unique ID used to multiplex responses on the client side
///
/// Treated as an opaque value on the server that is returned back to the client with the response
/// to a request.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ClientId(usize);

/// Spawns the server, and manages the IPC connection
///
/// Messages are sent though IPC and responses are dispatched back to the correct client based on
/// the received client ID.
#[derive(Debug)]
struct ClientDispatcher {
    /// The server task/process
    ///
    /// When dropped, this will block until the server process has quit. This field is explicitly
    /// owned by this struct and not reference counted in order to guarantee that this happens.
    server: RendererServer,

    /// The connection to the server task/process
    ///
    /// This will no longer send messages after the server process has terminated.
    conn: Arc<ClientConnection>,

    /// Each `ClientId` indexes into this field
    ///
    /// Using `RwLock` allows sending multiple times concurrently using `read()` and also allows
    /// more clients to be added using `write()`.
    clients: Arc<RwLock<Vec<mpsc::UnboundedSender<Result<ServerResponse, Disconnected>>>>>,
}

impl ClientDispatcher {
    async fn new() -> Result<Self, ConnectionError> {
        let (server, conn) = RendererServer::spawn().await?;
        let conn = Arc::new(conn);
        let clients = Arc::new(RwLock::new(Vec::<mpsc::UnboundedSender<_>>::new()));

        let task_conn = conn.clone();
        let task_clients = clients.clone();
        tokio::spawn(async move {
            loop {
                let (id, response) = match task_conn.recv().await {
                    Ok((id, response)) => (id, Ok(response)),

                    Err(IpcError::Disconnected) => {
                        // Alert all the clients of the disconnection
                        let clients = task_clients.read().await;
                        for client in &*clients {
                            // Ignoring the error since it just means that this particular client
                            // connection must have gotten dropped
                            client.send(Err(Disconnected)).unwrap_or(());
                        }
                        break;
                    },

                    Err(err) => panic!("Error while receiving IPC message: {:?}", err),
                };

                let clients = task_clients.read().await;

                let ClientId(index) = id;
                // Ignoring the error since it just means that this particular client
                // connection must have gotten dropped
                clients[index].send(response).unwrap_or(());
            }
        });

        Ok(Self {server, conn, clients})
    }

    async fn add_client(&self) -> (ClientId, mpsc::UnboundedReceiver<Result<ServerResponse, Disconnected>>) {
        let mut clients = self.clients.write().await;

        let id = ClientId(clients.len());
        let (sender, receiver) = mpsc::unbounded_channel();
        clients.push(sender);

        (id, receiver)
    }

    async fn send(&self, id: ClientId, req: ClientRequest) -> Result<(), ipc_channel::Error> {
        self.conn.send(id, req).await
    }
}

/// Represents a single connection to the server
#[derive(Debug)]
pub struct RendererClient {
    dispatcher: Arc<ClientDispatcher>,
    id: ClientId,
    receiver: Mutex<mpsc::UnboundedReceiver<Result<ServerResponse, Disconnected>>>,
}

impl RendererClient {
    /// Spawns a new server process and creates a connection to it
    pub async fn new() -> Result<Self, ConnectionError> {
        let dispatcher = Arc::new(ClientDispatcher::new().await?);
        let (id, receiver) = dispatcher.add_client().await;
        let receiver = Mutex::new(receiver);

        Ok(Self {dispatcher, id, receiver})
    }

    /// Creates a new renderer client that can also communicate to the same server
    pub async fn split(&self) -> Self {
        let dispatcher = self.dispatcher.clone();
        let (id, receiver) = dispatcher.add_client().await;
        let receiver = Mutex::new(receiver);

        Self {dispatcher, id, receiver}
    }

    /// Sends a message to the server process
    ///
    /// When possible, prefer using methods from `ProtocolClient` instead of using this directly
    pub async fn send(&self, req: ClientRequest) {
        // The error produced by send is a serialization error, so it signals a bug in this code,
        // not something that should be propagated to be handled elsewhere.
        self.dispatcher.send(self.id, req).await
            .expect("bug: error while sending message through IPC")
    }

    /// Receives a response from the server process
    ///
    /// When possible, prefer using methods from `ProtocolClient` instead of using this directly
    pub async fn recv(&self) -> ServerResponse {
        let mut receiver = self.receiver.lock().await;
        receiver.recv().await
            // Since this struct keeps a ref-counted copy of the senders, they can't have possibly
            // been dropped at this point.
            .expect("bug: client senders should not be dropped yet")
            // This panic causes the program to exit if turtle commands continue after the window
            // closes
            .unwrap_or_else(|err| panic!("IPC response not received: {}", err))
    }
}
<commit_msg>Realized that "dispatch" was a more appropriate word than "multiplex"<commit_after>use std::sync::Arc;

use ipc_channel::ipc::IpcError;
use serde::{Serialize, Deserialize};
use tokio::sync::{mpsc, RwLock, Mutex};
use thiserror::Error;

use crate::ipc_protocol::{ClientConnection, ConnectionError, ClientRequest, ServerResponse};
use crate::renderer_server::RendererServer;

/// Signals that the IPC connection has been disconnected and therefore the window was probably
/// closed
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
#[error("Cannot continue to run turtle commands after window is closed. This panic stops the thread, but is not necessarily an error.")]
struct Disconnected;

/// A unique ID used to dispatch responses on the client side
///
/// Treated as an opaque value on the server that is returned back to the client with the response
/// to a request.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ClientId(usize);

/// Spawns the server, and manages the IPC connection
///
/// Messages are sent though IPC and responses are dispatched back to the correct client based on
/// the received client ID.
#[derive(Debug)]
struct ClientDispatcher {
    /// The server task/process
    ///
    /// When dropped, this will block until the server process has quit. This field is explicitly
    /// owned by this struct and not reference counted in order to guarantee that this happens.
    server: RendererServer,

    /// The connection to the server task/process
    ///
    /// This will no longer send messages after the server process has terminated.
    conn: Arc<ClientConnection>,

    /// Each `ClientId` indexes into this field
    ///
    /// Using `RwLock` allows sending multiple times concurrently using `read()` and also allows
    /// more clients to be added using `write()`.
    clients: Arc<RwLock<Vec<mpsc::UnboundedSender<Result<ServerResponse, Disconnected>>>>>,
}

impl ClientDispatcher {
    async fn new() -> Result<Self, ConnectionError> {
        let (server, conn) = RendererServer::spawn().await?;
        let conn = Arc::new(conn);
        let clients = Arc::new(RwLock::new(Vec::<mpsc::UnboundedSender<_>>::new()));

        let task_conn = conn.clone();
        let task_clients = clients.clone();
        tokio::spawn(async move {
            loop {
                let (id, response) = match task_conn.recv().await {
                    Ok((id, response)) => (id, Ok(response)),

                    Err(IpcError::Disconnected) => {
                        // Alert all the clients of the disconnection
                        let clients = task_clients.read().await;
                        for client in &*clients {
                            // Ignoring the error since it just means that this particular client
                            // connection must have gotten dropped
                            client.send(Err(Disconnected)).unwrap_or(());
                        }
                        break;
                    },

                    Err(err) => panic!("Error while receiving IPC message: {:?}", err),
                };

                let clients = task_clients.read().await;

                let ClientId(index) = id;
                // Ignoring the error since it just means that this particular client
                // connection must have gotten dropped
                clients[index].send(response).unwrap_or(());
            }
        });

        Ok(Self {server, conn, clients})
    }

    async fn add_client(&self) -> (ClientId, mpsc::UnboundedReceiver<Result<ServerResponse, Disconnected>>) {
        let mut clients = self.clients.write().await;

        let id = ClientId(clients.len());
        let (sender, receiver) = mpsc::unbounded_channel();
        clients.push(sender);

        (id, receiver)
    }

    async fn send(&self, id: ClientId, req: ClientRequest) -> Result<(), ipc_channel::Error> {
        self.conn.send(id, req).await
    }
}

/// Represents a single connection to the server
#[derive(Debug)]
pub struct RendererClient {
    dispatcher: Arc<ClientDispatcher>,
    id: ClientId,
    receiver: Mutex<mpsc::UnboundedReceiver<Result<ServerResponse, Disconnected>>>,
}

impl RendererClient {
    /// Spawns a new server process and creates a connection to it
    pub async fn new() -> Result<Self, ConnectionError> {
        let dispatcher = Arc::new(ClientDispatcher::new().await?);
        let (id, receiver) = dispatcher.add_client().await;
        let receiver = Mutex::new(receiver);

        Ok(Self {dispatcher, id, receiver})
    }

    /// Creates a new renderer client that can also communicate to the same server
    pub async fn split(&self) -> Self {
        let dispatcher = self.dispatcher.clone();
        let (id, receiver) = dispatcher.add_client().await;
        let receiver = Mutex::new(receiver);

        Self {dispatcher, id, receiver}
    }

    /// Sends a message to the server process
    ///
    /// When possible, prefer using methods from `ProtocolClient` instead of using this directly
    pub async fn send(&self, req: ClientRequest) {
        // The error produced by send is a serialization error, so it signals a bug in this code,
        // not something that should be propagated to be handled elsewhere.
        self.dispatcher.send(self.id, req).await
            .expect("bug: error while sending message through IPC")
    }

    /// Receives a response from the server process
    ///
    /// When possible, prefer using methods from `ProtocolClient` instead of using this directly
    pub async fn recv(&self) -> ServerResponse {
        let mut receiver = self.receiver.lock().await;
        receiver.recv().await
            // Since this struct keeps a ref-counted copy of the senders, they can't have possibly
            // been dropped at this point.
            .expect("bug: client senders should not be dropped yet")
            // This panic causes the program to exit if turtle commands continue after the window
            // closes
            .unwrap_or_else(|err| panic!("IPC response not received: {}", err))
    }
}
<|endoftext|>