<commit_before>use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use std::thread;

use crossbeam_deque::{self as deque, Pop, Steal, Stealer, Worker};
use rand::thread_rng;
use rand::distributions::{Distribution, Uniform};
use threadpool::ThreadPool;

use gc::root::Slot;
use gc::{Address, Region};

pub fn start(rootset: &[Slot], heap: Region, perm: Region, number_workers: usize) {
    let mut workers = Vec::with_capacity(number_workers);
    let mut stealers = Vec::with_capacity(number_workers);

    for _ in 0..number_workers {
        let (w, s) = deque::lifo();
        workers.push(w);
        stealers.push(s);
    }

    for root in rootset {
        let root_ptr = root.get();

        if heap.contains(root_ptr) {
            let root_obj = root_ptr.to_mut_obj();

            if !root_obj.header().is_marked_non_atomic() {
                root_obj.header_mut().mark_non_atomic();
                workers[0].push(root_ptr);
            }
        } else {
            debug_assert!(root_ptr.is_null() || perm.contains(root_ptr));
        }
    }

    let pool = ThreadPool::with_name("gc-worker".to_string(), number_workers);
    let nworkers_stage = Arc::new(AtomicUsize::new(number_workers));

    for (task_id, worker) in workers.into_iter().enumerate() {
        let heap_region = heap.clone();
        let perm_region = perm.clone();

        let stealers = stealers.clone();
        let nworkers_stage = nworkers_stage.clone();

        pool.execute(move || {
            let mut local_segment = Segment::new();

            loop {
                let object_addr = if !local_segment.is_empty() {
                    local_segment.pop().expect("should be non-empty")
                } else {
                    if let Some(address) = pop(task_id, &worker, &stealers) {
                        address
                    } else if try_terminate(&nworkers_stage) {
                        break;
                    } else {
                        continue;
                    }
                };

                let object = object_addr.to_mut_obj();

                object.visit_reference_fields(|field| {
                    trace_slot(field, &mut local_segment, &worker, heap_region.clone(), perm_region.clone());
                });
            }
        });
    }

    pool.join();
}

fn trace_slot(slot: Slot, local_segment: &mut Segment, worker: &Worker<Address>, heap_region: Region, perm_region: Region) {
    let field_addr = slot.get();

    if heap_region.contains(field_addr) {
        let field_obj = field_addr.to_mut_obj();

        if field_obj.header().try_mark_non_atomic() {
            if local_segment.has_capacity() {
                local_segment.push(field_addr);
            } else {
                worker.push(field_addr);
            }
        }
    } else {
        debug_assert!(field_addr.is_null() || perm_region.contains(field_addr));
    }
}

fn pop(task_id: usize, worker: &Worker<Address>, stealers: &[Stealer<Address>]) -> Option<Address> {
    loop {
        match worker.pop() {
            Pop::Empty => break,
            Pop::Data(address) => {
                return Some(address)
            }
            Pop::Retry => continue,
        }
    }

    if stealers.len() == 1 {
        return None;
    }

    let mut rng = thread_rng();
    let range = Uniform::new(0, stealers.len());

    for _ in 0 .. 3 * stealers.len() {
        let mut stealer_id = task_id;

        while stealer_id == task_id {
            stealer_id = range.sample(&mut rng);
        }

        let stealer = &stealers[stealer_id];

        loop {
            match stealer.steal() {
                Steal::Empty => break,
                Steal::Data(address) => {
                    return Some(address)
                }
                Steal::Retry => continue,
            }
        }
    }

    None
}

fn try_terminate(nworkers_stage: &AtomicUsize) -> bool {
    enter_stage(nworkers_stage);
    thread::sleep(Duration::from_millis(1));
    !try_exit(nworkers_stage)
}

fn enter_stage(atomic: &AtomicUsize) -> bool {
    atomic.fetch_sub(1, Ordering::SeqCst) == 1
}

fn try_exit(atomic: &AtomicUsize) -> bool {
    let mut nworkers = atomic.load(Ordering::Relaxed);

    loop {
        if nworkers == 0 {
            return false;
        }
        let prev_nworkers = atomic.compare_and_swap(nworkers, nworkers + 1, Ordering::SeqCst);

        if nworkers == prev_nworkers {
            return true;
        }

        nworkers = prev_nworkers;
    }
}

const SEGMENT_SIZE: usize = 8;

struct Segment {
    data: Vec<Address>,
}

impl Segment {
    fn new() -> Segment {
        Segment {
            data: Vec::with_capacity(SEGMENT_SIZE),
        }
    }

    fn empty() -> Segment {
        Segment { data: Vec::new() }
    }

    fn with(addr: Address) -> Segment {
        let mut segment = Segment::new();
        segment.data.push(addr);

        segment
    }

    fn has_capacity(&self) -> bool {
        self.data.len() < SEGMENT_SIZE
    }

    fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    fn push(&mut self, addr: Address) {
        debug_assert!(self.has_capacity());
        self.data.push(addr);
    }

    fn pop(&mut self) -> Option<Address> {
        self.data.pop()
    }
}
<commit_msg>swiper: decrease pause time<commit_after>use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use std::thread;

use crossbeam_deque::{self as deque, Pop, Steal, Stealer, Worker};
use rand::thread_rng;
use rand::distributions::{Distribution, Uniform};
use threadpool::ThreadPool;

use gc::root::Slot;
use gc::{Address, Region};

pub fn start(rootset: &[Slot], heap: Region, perm: Region, number_workers: usize) {
    let mut workers = Vec::with_capacity(number_workers);
    let mut stealers = Vec::with_capacity(number_workers);

    for _ in 0..number_workers {
        let (w, s) = deque::lifo();
        workers.push(w);
        stealers.push(s);
    }

    for root in rootset {
        let root_ptr = root.get();

        if heap.contains(root_ptr) {
            let root_obj = root_ptr.to_mut_obj();

            if !root_obj.header().is_marked_non_atomic() {
                root_obj.header_mut().mark_non_atomic();
                workers[0].push(root_ptr);
            }
        } else {
            debug_assert!(root_ptr.is_null() || perm.contains(root_ptr));
        }
    }

    let pool = ThreadPool::with_name("gc-worker".to_string(), number_workers);
    let nworkers_stage = Arc::new(AtomicUsize::new(number_workers));

    for (task_id, worker) in workers.into_iter().enumerate() {
        let heap_region = heap.clone();
        let perm_region = perm.clone();

        let stealers = stealers.clone();
        let nworkers_stage = nworkers_stage.clone();

        pool.execute(move || {
            let mut local_segment = Segment::new();

            loop {
                let object_addr = if !local_segment.is_empty() {
                    local_segment.pop().expect("should be non-empty")
                } else {
                    if let Some(address) = pop(task_id, &worker, &stealers) {
                        address
                    } else if try_terminate(&nworkers_stage) {
                        break;
                    } else {
                        continue;
                    }
                };

                let object = object_addr.to_mut_obj();

                object.visit_reference_fields(|field| {
                    trace_slot(field, &mut local_segment, &worker, heap_region.clone(), perm_region.clone());
                });
            }
        });
    }

    pool.join();
}

fn trace_slot(slot: Slot, local_segment: &mut Segment, worker: &Worker<Address>, heap_region: Region, perm_region: Region) {
    let field_addr = slot.get();

    if heap_region.contains(field_addr) {
        let field_obj = field_addr.to_mut_obj();

        if field_obj.header().try_mark_non_atomic() {
            if local_segment.has_capacity() {
                local_segment.push(field_addr);
            } else {
                worker.push(field_addr);
            }
        }
    } else {
        debug_assert!(field_addr.is_null() || perm_region.contains(field_addr));
    }
}

fn pop(task_id: usize, worker: &Worker<Address>, stealers: &[Stealer<Address>]) -> Option<Address> {
    loop {
        match worker.pop() {
            Pop::Empty => break,
            Pop::Data(address) => {
                return Some(address)
            }
            Pop::Retry => continue,
        }
    }

    if stealers.len() == 1 {
        return None;
    }

    let mut rng = thread_rng();
    let range = Uniform::new(0, stealers.len());

    for _ in 0 .. 2 * stealers.len() {
        let mut stealer_id = task_id;

        while stealer_id == task_id {
            stealer_id = range.sample(&mut rng);
        }

        let stealer = &stealers[stealer_id];

        loop {
            match stealer.steal() {
                Steal::Empty => break,
                Steal::Data(address) => {
                    return Some(address)
                }
                Steal::Retry => continue,
            }
        }
    }

    None
}

fn try_terminate(nworkers_stage: &AtomicUsize) -> bool {
    enter_stage(nworkers_stage);
    thread::sleep(Duration::from_micros(1));
    !try_exit(nworkers_stage)
}

fn enter_stage(atomic: &AtomicUsize) -> bool {
    atomic.fetch_sub(1, Ordering::SeqCst) == 1
}

fn try_exit(atomic: &AtomicUsize) -> bool {
    let mut nworkers = atomic.load(Ordering::Relaxed);

    loop {
        if nworkers == 0 {
            return false;
        }
        let prev_nworkers = atomic.compare_and_swap(nworkers, nworkers + 1, Ordering::SeqCst);

        if nworkers == prev_nworkers {
            return true;
        }

        nworkers = prev_nworkers;
    }
}

const SEGMENT_SIZE: usize = 8;

struct Segment {
    data: Vec<Address>,
}

impl Segment {
    fn new() -> Segment {
        Segment {
            data: Vec::with_capacity(SEGMENT_SIZE),
        }
    }

    fn empty() -> Segment {
        Segment { data: Vec::new() }
    }

    fn with(addr: Address) -> Segment {
        let mut segment = Segment::new();
        segment.data.push(addr);

        segment
    }

    fn has_capacity(&self) -> bool {
        self.data.len() < SEGMENT_SIZE
    }

    fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    fn push(&mut self, addr: Address) {
        debug_assert!(self.has_capacity());
        self.data.push(addr);
    }

    fn pop(&mut self) -> Option<Address> {
        self.data.pop()
    }
}
<|endoftext|>


<commit_before>// -*- rust -*-
// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// xfail-test
type rec = {f: int};
fn f(p: *rec) -> int {

    // Test that * ptrs do not autoderef.  There is a deeper reason for
    // prohibiting this, beyond making unsafe things annoying (which doesn't
    // actually seem desirable to me).  The deeper reason is that if you
    // have a type like:
    //
    //    enum foo = *foo;
    //
    // you end up with an infinite auto-deref chain, which is
    // currently impossible (in all other cases, infinite auto-derefs
    // are prohibited by various checks, such as that the enum is
    // instantiable and so forth).

    return p.f; //~ ERROR attempted access of field `f` on type `*rec`
}

fn main() {
}
<commit_msg>testsuite: Update and un-xfail unsafe-fn-autoderef<commit_after>// -*- rust -*-
// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

struct Rec {
    f: int
}

fn f(p: *Rec) -> int {

    // Test that * ptrs do not autoderef.  There is a deeper reason for
    // prohibiting this, beyond making unsafe things annoying (which doesn't
    // actually seem desirable to me).  The deeper reason is that if you
    // have a type like:
    //
    //    enum foo = *foo;
    //
    // you end up with an infinite auto-deref chain, which is
    // currently impossible (in all other cases, infinite auto-derefs
    // are prohibited by various checks, such as that the enum is
    // instantiable and so forth).

    return p.f; //~ ERROR attempted access of field `f` on type `*Rec`
}

fn main() {
}
<|endoftext|>


<commit_before>extern crate simple_csv;

use std::env;
use std::fs::File;
use std::io::BufReader;
use std::str::FromStr;
use simple_csv::{SimpleCsvWriter, SimpleCsvReader};

fn main(){
    let args : Vec<String> = env::args().collect();
    let alpha = f64::from_str(&args[1]).expect("first argument should be the threshold");

    let f = File::open("assets/filter.csv").unwrap();
    let buf = BufReader::new(f);
    let reader = SimpleCsvReader::new(buf);

    let raw: Vec<(f64, f64)> = reader
        .map (|r| r.unwrap())
        .map(data)
        .collect();

    let mut sequence: Vec<DetrendData> = vec!();
    let mut data: Option<DetrendData> = None;
    for candidate in raw {
        match data{
            Some(previous) => {
                let next = previous.next(candidate, alpha);
                sequence.push(previous);
                data = Some(next);
            }

            None => {
                data = Some(DetrendData::initial(candidate))

            }
        }
    }

    let o = File::create("assets/detrend.csv").unwrap();
    let mut writer = SimpleCsvWriter::new(o);

    for data in sequence {
        writer.write(
            &data.as_vec()
        ).unwrap();
    }
}

fn data(row: Vec<String>) -> (f64, f64) {
    let iter = row.iter();

    let raw: Vec<f64> = iter
        .map(|s| f64::from_str(s).unwrap())
        .collect();

    (raw[0], raw[1])
}

struct DetrendData {
    time: f64,
    brightness: f64,
    trend: f64,
    difference: f64,
}

impl DetrendData {
    fn initial((time, brightness): (f64,f64)) -> DetrendData {
        DetrendData {
            time: time,
            brightness: brightness,
            trend: brightness,
            difference: 0f64,
        }
    }

    fn next(&self, (time, brightness): (f64, f64), alpha: f64) -> DetrendData {
        let trend = alpha * brightness + (1f64 - alpha) * self.trend;
        DetrendData {
            time: time,
            brightness: brightness,
            trend: trend,
            difference: brightness - trend,
        }

    }

    fn as_vec(&self) -> Vec<String> {
        vec!(
            self.time.to_string(),
            self.brightness.to_string(),
            self.trend.to_string(),
            self.difference.to_string(),
        )
    }
}
<commit_msg>Determine trend from brightness<commit_after>extern crate simple_csv;

use std::env;
use std::fs::File;
use std::io::BufReader;
use std::str::FromStr;
use simple_csv::{SimpleCsvWriter, SimpleCsvReader};

fn main(){
    let args : Vec<String> = env::args().collect();
    let alpha = f64::from_str(&args[1]).expect("first argument should be the threshold");

    let f = File::open("assets/brightness.csv").unwrap();
    let buf = BufReader::new(f);
    let reader = SimpleCsvReader::new(buf);

    let raw: Vec<(f64, f64)> = reader
        .map (|r| r.unwrap())
        .map(data)
        .collect();

    let mut sequence: Vec<DetrendData> = vec!();
    let mut data: Option<DetrendData> = None;
    for candidate in raw {
        match data{
            Some(previous) => {
                let next = previous.next(candidate, alpha);
                sequence.push(previous);
                data = Some(next);
            }

            None => {
                data = Some(DetrendData::initial(candidate))

            }
        }
    }

    let o = File::create("assets/detrend.csv").unwrap();
    let mut writer = SimpleCsvWriter::new(o);

    for data in sequence {
        writer.write(
            &data.as_vec()
        ).unwrap();
    }
}

fn data(row: Vec<String>) -> (f64, f64) {
    let iter = row.iter();

    let raw: Vec<f64> = iter
        .map(|s| f64::from_str(s).unwrap())
        .collect();

    (raw[0], raw[1])
}

struct DetrendData {
    time: f64,
    brightness: f64,
    trend: f64,
}

impl DetrendData {
    fn initial((time, brightness): (f64,f64)) -> DetrendData {
        DetrendData {
            time: time,
            brightness: brightness,
            trend: brightness,
        }
    }

    fn next(&self, (time, brightness): (f64, f64), alpha: f64) -> DetrendData {
        let trend = alpha * brightness + (1f64 - alpha) * self.trend;
        DetrendData {
            time: time,
            brightness: brightness,
            trend: trend,
        }

    }

    fn as_vec(&self) -> Vec<String> {
        vec!(
            self.time.to_string(),
            self.brightness.to_string(),
            self.trend.to_string(),
        )
    }
}
<|endoftext|>


<commit_before>use crate::{command::studio::enter::ARTIFACT_PATH_ENVVAR,
            common::ui::UI,
            error::{Error,
                    Result},
            hcore::{crypto::default_cache_key_path,
                    env as henv,
                    fs::{find_command,
                         CACHE_ARTIFACT_PATH,
                         CACHE_KEY_PATH},
                    os::process,
                    package::target},
            license,
            VERSION};
use atty;
use std::{env,
          ffi::{OsStr,
                OsString},
          path::{Path,
                 PathBuf},
          process::{Command,
                    Stdio}};

const DOCKER_CMD: &str = "docker";
const DOCKER_CMD_ENVVAR: &str = "HAB_DOCKER_BINARY";

const DOCKER_IMAGE: &str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, args: &[OsString]) -> Result<()> {
    let mut args = args.to_vec();
    if args.get(0) == Some(&OsString::from("rm")) {
        return Err(Error::CannotRemoveDockerStudio);
    }

    let docker_cmd = find_docker_cmd()?;

    if is_image_present(&docker_cmd) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd)?;
    }

    let mnt_prefix = if is_serving_windows_containers(&docker_cmd) {
        "c:"
    } else {
        ""
    };
    let mut volumes = vec![format!("{}:{}{}",
                                   env::current_dir().unwrap().to_string_lossy(),
                                   mnt_prefix,
                                   "/src"),
                           format!("{}:{}/{}",
                                   default_cache_key_path(None).to_string_lossy(),
                                   mnt_prefix,
                                   CACHE_KEY_PATH),];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!("{}:{}/{}",
                             cache_artifact_path, mnt_prefix, CACHE_ARTIFACT_PATH));
    }
    if !is_serving_windows_containers(&docker_cmd)
       && (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let mut env_vars = vec![String::from("DEBUG"),
                            String::from("DO_CHECK"),
                            String::from("HAB_AUTH_TOKEN"),
                            String::from("HAB_BLDR_URL"),
                            String::from("HAB_BLDR_CHANNEL"),
                            String::from("HAB_NOCOLORING"),
                            String::from("HAB_LICENSE"),
                            String::from("HAB_ORIGIN"),
                            String::from("HAB_ORIGIN_KEYS"),
                            String::from("HAB_STUDIO_BACKLINE_PKG"),
                            String::from("HAB_STUDIO_NOSTUDIORC"),
                            String::from("HAB_STUDIO_SUP"),
                            String::from("HAB_UPDATE_STRATEGY_FREQUENCY_MS"),
                            String::from("http_proxy"),
                            String::from("https_proxy"),
                            String::from("RUST_LOG"),];

    for (key, _) in env::vars() {
        if key.starts_with(HAB_STUDIO_SECRET) {
            env_vars.push(key);
        }
    }

    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }

    // Windows containers do not use filesystem sharing for
    // local mounts
    if !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    run_container(docker_cmd, &args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or_else(|_| DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    let image = image_identifier_for_active_target(&docker_cmd);

    cmd.arg("images").arg(&image).arg("-q");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path) -> Result<()> {
    let image = image_identifier_for_active_target(&docker_cmd);
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
       .arg(&image)
       .stdout(Stdio::inherit())
       .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd.spawn()
                    .expect("Docker command failed to spawn")
                    .wait_with_output()
                    .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", &image);
    } else {
        debug!("Pulling Docker image '{}' failed with exit code: {:?}",
               &image, result.status);

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(image.to_string()));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(image.to_string()));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I) -> Result<()>
    where I: IntoIterator<Item = S>,
          S: AsRef<OsStr>
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    let image = image_identifier_for_active_target(&docker_cmd);

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image.into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd).args(&cmd_args)
                                                 .output()
                                                 .expect("docker failed to start");

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty()
       && (stderr.as_str().contains("Mounts denied")
           || stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn run_container<I, J, S, T>(docker_cmd: PathBuf,
                             args: &[OsString],
                             volumes: I,
                             env_vars: J)
                             -> Result<()>
    where I: IntoIterator<Item = S>,
          J: IntoIterator<Item = T>,
          S: AsRef<OsStr>,
          T: AsRef<str>
{
    let using_windows_containers = is_serving_windows_containers(&docker_cmd);
    let image = image_identifier_for_active_target(&docker_cmd);
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];

    if !using_windows_containers {
        cmd_args.push("--privileged".into());
    }

    if atty::is(atty::Stream::Stderr) || atty::is(atty::Stream::Stdout) {
        cmd_args.push("--tty".into());
        cmd_args.push("--interactive".into());
    }

    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts
            .split(' ')
            .map(|v| v.into())
            // Ensure we're not passing something like `--tty` again here.
            .filter(|v| !cmd_args.contains(v))
            .collect::<Vec<_>>();

        if !opts.is_empty() {
            debug!("Adding extra Docker options from {} = {:?}",
                   DOCKER_OPTS_ENVVAR, opts);
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }

    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        } else if var.as_ref() == "HAB_LICENSE" && license::license_exists() {
            debug!("Hab license already accepted. Setting container env var: \
                    HAB_LICENSE=accept-no-persist");
            cmd_args.push("--env".into());
            cmd_args.push("HAB_LICENSE=accept-no-persist".to_string().into());
        }
    }

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }

    cmd_args.push(image.into());
    cmd_args.extend_from_slice(args);

    if using_windows_containers {
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }

    unset_proxy_env_vars();
    process::become_command(docker_cmd, &cmd_args)?;
    Ok(())
}

fn unset_proxy_env_vars() {
    for var in &["http_proxy", "https_proxy"] {
        if henv::var(var).is_ok() {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

fn image_identifier_for_active_target(docker_cmd: &Path) -> String {
    image_identifier(is_serving_windows_containers(docker_cmd),
                     target::PackageTarget::active_target())
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(using_windows_containers: bool, target: target::PackageTarget) -> String {
    let version: Vec<&str> = VERSION.split('/').collect();
    let (img, studio_target) = if using_windows_containers {
        (DOCKER_WINDOWS_IMAGE, target::X86_64_WINDOWS)
    } else {
        let t = match target {
            target::X86_64_LINUX_KERNEL2 => target::X86_64_LINUX_KERNEL2,
            _ => target::X86_64_LINUX,
        };
        (DOCKER_IMAGE, t)
    };

    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or_else(|_| {
                                      format!("{}-{}:{}", img, studio_target, version[0])
                                  })
}

#[cfg(test)]
mod tests {
    use super::{image_identifier,
                DOCKER_IMAGE,
                DOCKER_WINDOWS_IMAGE};
    use crate::VERSION;

    use crate::hcore::package::target;

    #[test]
    fn retrieve_image_identifier() {
        let windows_container = true;
        assert_eq!(image_identifier(!windows_container, target::X86_64_DARWIN),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_LINUX_KERNEL2),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux-kernel2", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(windows_container, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION));
        assert_eq!(image_identifier(windows_container, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION));
    }
}
<commit_msg>Checks if default cache key path exists and tells user if it doesn't.<commit_after>use crate::{command::studio::enter::ARTIFACT_PATH_ENVVAR,
            common::ui::UI,
            error::{Error,
                    Result},
            hcore::{crypto::default_cache_key_path,
                    env as henv,
                    fs::{find_command,
                         CACHE_ARTIFACT_PATH,
                         CACHE_KEY_PATH},
                    os::process,
                    package::target},
            license,
            VERSION};
use atty;
use std::{env,
          ffi::{OsStr,
                OsString},
          path::{Path,
                 PathBuf},
          process::{Command,
                    Stdio}};

const DOCKER_CMD: &str = "docker";
const DOCKER_CMD_ENVVAR: &str = "HAB_DOCKER_BINARY";

const DOCKER_IMAGE: &str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, args: &[OsString]) -> Result<()> {
    let mut args = args.to_vec();
    if args.get(0) == Some(&OsString::from("rm")) {
        return Err(Error::CannotRemoveDockerStudio);
    }

    let docker_cmd = find_docker_cmd()?;

    if is_image_present(&docker_cmd) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd)?;
    }

    let mnt_prefix = if is_serving_windows_containers(&docker_cmd) {
        "c:"
    } else {
        ""
    };

    let local_cache_key_path = default_cache_key_path(None);
    if !local_cache_key_path.exists() {
        return Err(Error::FileNotFound(
            format!(
                "{}\nRun `hab setup` to create an origin or \
                use `hab origin key` to configure your keys.",
                local_cache_key_path.to_string_lossy().into_owned()
            )
        ));
    }

    let mut volumes = vec![format!("{}:{}{}",
                                   env::current_dir().unwrap().to_string_lossy(),
                                   mnt_prefix,
                                   "/src"),
                           format!("{}:{}/{}",
                                   local_cache_key_path.to_string_lossy(),
                                   mnt_prefix,
                                   CACHE_KEY_PATH),];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!("{}:{}/{}",
                             cache_artifact_path, mnt_prefix, CACHE_ARTIFACT_PATH));
    }
    if !is_serving_windows_containers(&docker_cmd)
       && (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let mut env_vars = vec![String::from("DEBUG"),
                            String::from("DO_CHECK"),
                            String::from("HAB_AUTH_TOKEN"),
                            String::from("HAB_BLDR_URL"),
                            String::from("HAB_BLDR_CHANNEL"),
                            String::from("HAB_NOCOLORING"),
                            String::from("HAB_LICENSE"),
                            String::from("HAB_ORIGIN"),
                            String::from("HAB_ORIGIN_KEYS"),
                            String::from("HAB_STUDIO_BACKLINE_PKG"),
                            String::from("HAB_STUDIO_NOSTUDIORC"),
                            String::from("HAB_STUDIO_SUP"),
                            String::from("HAB_UPDATE_STRATEGY_FREQUENCY_MS"),
                            String::from("http_proxy"),
                            String::from("https_proxy"),
                            String::from("RUST_LOG"),];

    for (key, _) in env::vars() {
        if key.starts_with(HAB_STUDIO_SECRET) {
            env_vars.push(key);
        }
    }

    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }

    // Windows containers do not use filesystem sharing for
    // local mounts
    if !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    run_container(docker_cmd, &args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or_else(|_| DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    let image = image_identifier_for_active_target(&docker_cmd);

    cmd.arg("images").arg(&image).arg("-q");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path) -> Result<()> {
    let image = image_identifier_for_active_target(&docker_cmd);
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
       .arg(&image)
       .stdout(Stdio::inherit())
       .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd.spawn()
                    .expect("Docker command failed to spawn")
                    .wait_with_output()
                    .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", &image);
    } else {
        debug!("Pulling Docker image '{}' failed with exit code: {:?}",
               &image, result.status);

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(image.to_string()));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(image.to_string()));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I) -> Result<()>
    where I: IntoIterator<Item = S>,
          S: AsRef<OsStr>
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    let image = image_identifier_for_active_target(&docker_cmd);

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image.into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd).args(&cmd_args)
                                                 .output()
                                                 .expect("docker failed to start");

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty()
       && (stderr.as_str().contains("Mounts denied")
           || stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn run_container<I, J, S, T>(docker_cmd: PathBuf,
                             args: &[OsString],
                             volumes: I,
                             env_vars: J)
                             -> Result<()>
    where I: IntoIterator<Item = S>,
          J: IntoIterator<Item = T>,
          S: AsRef<OsStr>,
          T: AsRef<str>
{
    let using_windows_containers = is_serving_windows_containers(&docker_cmd);
    let image = image_identifier_for_active_target(&docker_cmd);
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];

    if !using_windows_containers {
        cmd_args.push("--privileged".into());
    }

    if atty::is(atty::Stream::Stderr) || atty::is(atty::Stream::Stdout) {
        cmd_args.push("--tty".into());
        cmd_args.push("--interactive".into());
    }

    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts
            .split(' ')
            .map(|v| v.into())
            // Ensure we're not passing something like `--tty` again here.
            .filter(|v| !cmd_args.contains(v))
            .collect::<Vec<_>>();

        if !opts.is_empty() {
            debug!("Adding extra Docker options from {} = {:?}",
                   DOCKER_OPTS_ENVVAR, opts);
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }

    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        } else if var.as_ref() == "HAB_LICENSE" && license::license_exists() {
            debug!("Hab license already accepted. Setting container env var: \
                    HAB_LICENSE=accept-no-persist");
            cmd_args.push("--env".into());
            cmd_args.push("HAB_LICENSE=accept-no-persist".to_string().into());
        }
    }

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }

    cmd_args.push(image.into());
    cmd_args.extend_from_slice(args);

    if using_windows_containers {
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }

    unset_proxy_env_vars();
    process::become_command(docker_cmd, &cmd_args)?;
    Ok(())
}

fn unset_proxy_env_vars() {
    for var in &["http_proxy", "https_proxy"] {
        if henv::var(var).is_ok() {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

fn image_identifier_for_active_target(docker_cmd: &Path) -> String {
    image_identifier(is_serving_windows_containers(docker_cmd),
                     target::PackageTarget::active_target())
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(using_windows_containers: bool, target: target::PackageTarget) -> String {
    let version: Vec<&str> = VERSION.split('/').collect();
    let (img, studio_target) = if using_windows_containers {
        (DOCKER_WINDOWS_IMAGE, target::X86_64_WINDOWS)
    } else {
        let t = match target {
            target::X86_64_LINUX_KERNEL2 => target::X86_64_LINUX_KERNEL2,
            _ => target::X86_64_LINUX,
        };
        (DOCKER_IMAGE, t)
    };

    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or_else(|_| {
                                      format!("{}-{}:{}", img, studio_target, version[0])
                                  })
}

#[cfg(test)]
mod tests {
    use super::{image_identifier,
                DOCKER_IMAGE,
                DOCKER_WINDOWS_IMAGE};
    use crate::VERSION;

    use crate::hcore::package::target;

    #[test]
    fn retrieve_image_identifier() {
        let windows_container = true;
        assert_eq!(image_identifier(!windows_container, target::X86_64_DARWIN),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_LINUX_KERNEL2),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux-kernel2", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(windows_container, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION));
        assert_eq!(image_identifier(windows_container, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION));
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2020 Jason Ish
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
// IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
// OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
// OR OTHER DEALINGS IN THE SOFTWARE.

extern crate nom;

use std::io::BufRead;

use nom::{
    bytes::complete::tag,
    character::complete::multispace0,
    error::{ErrorKind, ParseError},
    multi::many0,
    sequence::{preceded, tuple},
    IResult,
};

static WHITESPACE: &str = " \t\r\n";

#[derive(Debug, PartialEq)]
enum InternalError<I> {
    NoListEnd,
    Nom(I, ErrorKind),
}

impl<I> ParseError<I> for InternalError<I> {
    fn from_error_kind(input: I, kind: ErrorKind) -> Self {
        InternalError::Nom(input, kind)
    }

    fn append(_input: I, _kind: ErrorKind, other: Self) -> Self {
        other
    }
}

#[derive(Debug, PartialEq, Default)]
pub struct RuleHeader {
    action: String,
    proto: String,
    src_addr: String,
    src_port: String,
    direction: String,
    dst_addr: String,
    dst_port: String,
}

#[derive(Debug, PartialEq)]
pub struct RuleOption {
    pub key: String,
    pub val: Option<String>,
}

#[derive(Debug)]
pub struct TokenizedRule {
    pub header: RuleHeader,
    pub options: Vec<RuleOption>,
    pub original: String,
}

fn parse_header_token(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    preceded(multispace0, nom::bytes::complete::is_not(WHITESPACE))(input)
}

/// Get the next rule option key (name).
///
/// This will return the next set of characters terminated by ; or : with all
/// leading whitespace removed.
fn get_option_key(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    preceded(multispace0, nom::bytes::complete::is_not(";:"))(input)
}

/// Get the next option value.
///
/// This will return the next sequence of characters terminated by ';',
/// while removing any instances of an escaped ';'.
fn get_option_value(input: &str) -> IResult<&str, String, InternalError<&str>> {
    let mut output = Vec::new();
    let mut escaped = false;
    let mut end = 0;
    for (i, c) in input.chars().enumerate() {
        end = i;
        if c == '\\' {
            escaped = true;
        } else if escaped {
            if c == ';' {
                output.push(c);
            } else {
                output.push('\\');
                output.push(c);
            }
            escaped = false;
        } else if c == ';' {
            // Eat the ';'.
            end += 1;
            break;
        } else {
            output.push(c);
        }
    }
    let (_, rem) = input.split_at(end);
    Ok((rem, output.into_iter().collect()))
}

/// Parse a list as a token.
///
/// Returns the complete list as a single token, handling nested
/// lists, and spaces between listen items.
///
/// As this doesn't use nom, we make it look like a nom parser,
/// and use nom error types that are close enough.
fn parse_list_token(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    // Peek into the input to make sure it looks like the start of a list.
    let input = preceded(
        multispace0,
        nom::combinator::peek(nom::bytes::complete::tag("[")),
    )(input)?
    .0;
    let mut in_list = 0;
    let mut end = 0;

    for (i, c) in input.chars().enumerate() {
        match c {
            '[' => in_list += 1,
            ']' => in_list -= 1,
            _ => {}
        }

        if in_list == 0 {
            end = i + 1;
            break;
        }
    }

    // If we got here and are still in the list, then we ran out of input.
    if in_list > 0 {
        return Err(nom::Err::Error(InternalError::NoListEnd));
    }

    let (list, rem) = input.split_at(end);
    Ok((rem, list))
}

fn parse_header(input: &str) -> IResult<&str, RuleHeader, InternalError<&str>> {
    let maybe_list = &nom::branch::alt((parse_list_token, parse_header_token));
    let (rem, (action, proto, src_addr, src_port, direction, dst_addr, dst_port)) = tuple((
        // Action.
        parse_header_token,
        // Proto.
        maybe_list,
        // Source address.
        maybe_list,
        // Source port.
        maybe_list,
        // Direction.
        parse_header_token,
        // Destination address.
        maybe_list,
        // Destination port.
        maybe_list,
    ))(input)?;
    Ok((
        rem,
        RuleHeader {
            action: String::from(action),
            proto: String::from(proto),
            src_addr: String::from(src_addr),
            src_port: String::from(src_port),
            direction: String::from(direction),
            dst_addr: String::from(dst_addr),
            dst_port: String::from(dst_port),
        },
    ))
}

fn parse_option(input: &str) -> IResult<&str, RuleOption, InternalError<&str>> {
    // Eat any leading space, then parse up to a : or a ;.
    let (input, _) = multispace0(input)?;
    let (input, key) = get_option_key(input)?;
    let (input, sep) = nom::character::complete::one_of(";:")(input)?;
    if sep == ';' {
        return Ok((
            input,
            RuleOption {
                key: String::from(key),
                val: None,
            },
        ));
    }

    // Drop any leading whitespace then parse up to ';', accounting for
    // escaped occurrences of ';'.
    let (input, _) = multispace0(input)?;
    let (input, val) = get_option_value(input)?;

    Ok((
        input,
        RuleOption {
            key: String::from(key),
            val: Some(strip_quotes(&val)),
        },
    ))
}

/// Remove quotes from a string, but preserve any escaped quotes.
fn strip_quotes(input: &str) -> String {
    let mut escaped = false;
    let mut out: Vec<char> = Vec::new();

    for c in input.chars() {
        if escaped {
            out.push(c);
            escaped = false;
        } else {
            match c {
                '"' => {}
                '\\' => {
                    escaped = true;
                }
                _ => {
                    out.push(c);
                }
            }
        }
    }

    out.iter().collect()
}

fn internal_parse_rule(input: &str) -> IResult<&str, TokenizedRule, InternalError<&str>> {
    let original = String::from(input);
    let (input, header) = parse_header(input)?;
    let (input, _) = preceded(multispace0, tag("("))(input)?;
    let (input, options) = many0(parse_option)(input)?;
    let (input, _) = preceded(multispace0, tag(")"))(input)?;

    Ok((
        input,
        TokenizedRule {
            header,
            options,
            original,
        },
    ))
}

pub fn parse_rule(input: &str) -> anyhow::Result<TokenizedRule> {
    match internal_parse_rule(input) {
        Ok((_, rule)) => return Ok(rule),
        Err(err) => return Err(anyhow::anyhow!(err.to_string())),
    }
}

/// Read the next rule from a reader.
///
/// This will actually return any line it reads, but will join together
/// multiline rules into a single line rule.
///
/// Ok(None) will be returned on EOF, and an error will be returned on read
/// error.
pub fn read_next_rule(input: &mut dyn BufRead) -> Result<Option<String>, std::io::Error> {
    let mut line = String::new();
    loop {
        let mut tmp = String::new();
        let n = input.read_line(&mut tmp)?;
        if n == 0 {
            return Ok(None);
        }

        let tmp = tmp.trim();

        if !tmp.ends_with('\\') {
            line.push_str(tmp);
            break;
        }

        line.push_str(&tmp[..tmp.len() - 1]);
    }
    Ok(Some(line))
}

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

    #[test]
    fn test_parse_token() {
        assert_eq!(parse_header_token("alert"), Ok(("", "alert")));
        assert_eq!(parse_header_token(" alert"), Ok(("", "alert")));
        assert_eq!(parse_header_token(" alert "), Ok((" ", "alert")));
        assert_eq!(parse_header_token("http_uri"), Ok(("", "http_uri")));
    }

    #[test]
    fn test_parse_quoted_string() {
        assert_eq!(
            strip_quotes(r#""some quoted \" string""#),
            r#"some quoted " string"#
        );
    }

    #[test]
    fn test_parse_list() {
        assert_eq!(parse_list_token("[1]"), Ok(("", "[1]")));
        assert_eq!(parse_list_token("[1,2,3]"), Ok(("", "[1,2,3]")));
        assert_eq!(parse_list_token(" [1,2,3]"), Ok(("", "[1,2,3]")));
        assert_eq!(
            parse_list_token(" [1,2,3,[a,b,c]]"),
            Ok(("", "[1,2,3,[a,b,c]]"))
        );

        assert!(parse_list_token("1,2,3]").is_err());
        assert!(parse_list_token("[1,2,3").is_err());

        assert!(parse_list_token("token").is_err());
    }

    #[test]
    fn test_parse_header() {
        let rule = parse_header("alert tcp any any -> any any");
        assert_eq!(
            rule,
            Ok((
                "",
                RuleHeader {
                    action: String::from("alert"),
                    proto: String::from("tcp"),
                    src_addr: String::from("any"),
                    src_port: String::from("any"),
                    direction: String::from("->"),
                    dst_addr: String::from("any"),
                    dst_port: String::from("any"),
                    ..Default::default()
                }
            ))
        );
    }

    #[test]
    fn test_parse_option_without_value() {
        assert_eq!(
            Ok((
                "",
                RuleOption {
                    key: String::from("http_uri"),
                    val: None,
                }
            )),
            parse_option("http_uri;")
        );
    }

    #[test]
    fn test_parse_option() {
        assert_eq!(
            parse_option("msg:value;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("value".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option("msg:value with spaces;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("value with spaces".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option("msg:terminated value with spaces;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("terminated value with spaces".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: an escaped \; terminant; next_option;"#),
            Ok((
                " next_option;",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("an escaped ; terminant".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: "A Quoted Message";"#),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some(r#"A Quoted Message"#.to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: "A Quoted Message with escaped \" quotes.";"#),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some(r#"A Quoted Message with escaped " quotes."#.to_string()),
                }
            ))
        );
    }

    #[test]
    fn test_parse_rule() {
        let rule = parse_rule(r#"alert ip any any -> any any (msg:"Some Message with a \" Quote"; metadata: key, val; sid:1; rev:1;)"#).unwrap();
        assert_eq!(rule.header.action, "alert");
    }

    #[test]
    fn test_tokenize_option_value() {
        assert_eq!(
            get_option_value("one; two"),
            Ok((" two", String::from("one")))
        );

        assert_eq!(
            get_option_value("one\\;; two"),
            Ok((" two", String::from("one;")))
        );

        let optval = r#""/^(?:[a-zA-Z0-9_%+])*(?:[\x2c\x22\x27\x28]|\x252[c278])/PRi";"#;
        let expected = r#""/^(?:[a-zA-Z0-9_%+])*(?:[\x2c\x22\x27\x28]|\x252[c278])/PRi""#;
        assert_eq!(get_option_value(optval), Ok(("", String::from(expected))));
    }

    #[test]
    fn test_parse_from_reader() {
        let input = r#"alert ip any any -> any any (\
            msg:"TEST RULE"; sid:1; rev:1;)"#;
        let mut reader = input.as_bytes();
        let next = read_next_rule(&mut reader).unwrap();
        assert_eq!(
            next,
            Some(r#"alert ip any any -> any any (msg:"TEST RULE"; sid:1; rev:1;)"#.to_string())
        );
        assert_eq!(read_next_rule(&mut reader).unwrap(), None);
    }
}
<commit_msg>rule parser: make header fields public<commit_after>// Copyright (c) 2020 Jason Ish
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
// IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
// OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
// OR OTHER DEALINGS IN THE SOFTWARE.

extern crate nom;

use std::io::BufRead;

use nom::{
    bytes::complete::tag,
    character::complete::multispace0,
    error::{ErrorKind, ParseError},
    multi::many0,
    sequence::{preceded, tuple},
    IResult,
};

static WHITESPACE: &str = " \t\r\n";

#[derive(Debug, PartialEq)]
enum InternalError<I> {
    NoListEnd,
    Nom(I, ErrorKind),
}

impl<I> ParseError<I> for InternalError<I> {
    fn from_error_kind(input: I, kind: ErrorKind) -> Self {
        InternalError::Nom(input, kind)
    }

    fn append(_input: I, _kind: ErrorKind, other: Self) -> Self {
        other
    }
}

#[derive(Debug, PartialEq, Default)]
pub struct RuleHeader {
    pub action: String,
    pub proto: String,
    pub src_addr: String,
    pub src_port: String,
    pub direction: String,
    pub dst_addr: String,
    pub dst_port: String,
}

#[derive(Debug, PartialEq)]
pub struct RuleOption {
    pub key: String,
    pub val: Option<String>,
}

#[derive(Debug)]
pub struct TokenizedRule {
    pub header: RuleHeader,
    pub options: Vec<RuleOption>,
    pub original: String,
}

fn parse_header_token(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    preceded(multispace0, nom::bytes::complete::is_not(WHITESPACE))(input)
}

/// Get the next rule option key (name).
///
/// This will return the next set of characters terminated by ; or : with all
/// leading whitespace removed.
fn get_option_key(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    preceded(multispace0, nom::bytes::complete::is_not(";:"))(input)
}

/// Get the next option value.
///
/// This will return the next sequence of characters terminated by ';',
/// while removing any instances of an escaped ';'.
fn get_option_value(input: &str) -> IResult<&str, String, InternalError<&str>> {
    let mut output = Vec::new();
    let mut escaped = false;
    let mut end = 0;
    for (i, c) in input.chars().enumerate() {
        end = i;
        if c == '\\' {
            escaped = true;
        } else if escaped {
            if c == ';' {
                output.push(c);
            } else {
                output.push('\\');
                output.push(c);
            }
            escaped = false;
        } else if c == ';' {
            // Eat the ';'.
            end += 1;
            break;
        } else {
            output.push(c);
        }
    }
    let (_, rem) = input.split_at(end);
    Ok((rem, output.into_iter().collect()))
}

/// Parse a list as a token.
///
/// Returns the complete list as a single token, handling nested
/// lists, and spaces between listen items.
///
/// As this doesn't use nom, we make it look like a nom parser,
/// and use nom error types that are close enough.
fn parse_list_token(input: &str) -> IResult<&str, &str, InternalError<&str>> {
    // Peek into the input to make sure it looks like the start of a list.
    let input = preceded(
        multispace0,
        nom::combinator::peek(nom::bytes::complete::tag("[")),
    )(input)?
    .0;
    let mut in_list = 0;
    let mut end = 0;

    for (i, c) in input.chars().enumerate() {
        match c {
            '[' => in_list += 1,
            ']' => in_list -= 1,
            _ => {}
        }

        if in_list == 0 {
            end = i + 1;
            break;
        }
    }

    // If we got here and are still in the list, then we ran out of input.
    if in_list > 0 {
        return Err(nom::Err::Error(InternalError::NoListEnd));
    }

    let (list, rem) = input.split_at(end);
    Ok((rem, list))
}

fn parse_header(input: &str) -> IResult<&str, RuleHeader, InternalError<&str>> {
    let maybe_list = &nom::branch::alt((parse_list_token, parse_header_token));
    let (rem, (action, proto, src_addr, src_port, direction, dst_addr, dst_port)) = tuple((
        // Action.
        parse_header_token,
        // Proto.
        maybe_list,
        // Source address.
        maybe_list,
        // Source port.
        maybe_list,
        // Direction.
        parse_header_token,
        // Destination address.
        maybe_list,
        // Destination port.
        maybe_list,
    ))(input)?;
    Ok((
        rem,
        RuleHeader {
            action: String::from(action),
            proto: String::from(proto),
            src_addr: String::from(src_addr),
            src_port: String::from(src_port),
            direction: String::from(direction),
            dst_addr: String::from(dst_addr),
            dst_port: String::from(dst_port),
        },
    ))
}

fn parse_option(input: &str) -> IResult<&str, RuleOption, InternalError<&str>> {
    // Eat any leading space, then parse up to a : or a ;.
    let (input, _) = multispace0(input)?;
    let (input, key) = get_option_key(input)?;
    let (input, sep) = nom::character::complete::one_of(";:")(input)?;
    if sep == ';' {
        return Ok((
            input,
            RuleOption {
                key: String::from(key),
                val: None,
            },
        ));
    }

    // Drop any leading whitespace then parse up to ';', accounting for
    // escaped occurrences of ';'.
    let (input, _) = multispace0(input)?;
    let (input, val) = get_option_value(input)?;

    Ok((
        input,
        RuleOption {
            key: String::from(key),
            val: Some(strip_quotes(&val)),
        },
    ))
}

/// Remove quotes from a string, but preserve any escaped quotes.
fn strip_quotes(input: &str) -> String {
    let mut escaped = false;
    let mut out: Vec<char> = Vec::new();

    for c in input.chars() {
        if escaped {
            out.push(c);
            escaped = false;
        } else {
            match c {
                '"' => {}
                '\\' => {
                    escaped = true;
                }
                _ => {
                    out.push(c);
                }
            }
        }
    }

    out.iter().collect()
}

fn internal_parse_rule(input: &str) -> IResult<&str, TokenizedRule, InternalError<&str>> {
    let original = String::from(input);
    let (input, header) = parse_header(input)?;
    let (input, _) = preceded(multispace0, tag("("))(input)?;
    let (input, options) = many0(parse_option)(input)?;
    let (input, _) = preceded(multispace0, tag(")"))(input)?;

    Ok((
        input,
        TokenizedRule {
            header,
            options,
            original,
        },
    ))
}

pub fn parse_rule(input: &str) -> anyhow::Result<TokenizedRule> {
    match internal_parse_rule(input) {
        Ok((_, rule)) => return Ok(rule),
        Err(err) => return Err(anyhow::anyhow!(err.to_string())),
    }
}

/// Read the next rule from a reader.
///
/// This will actually return any line it reads, but will join together
/// multiline rules into a single line rule.
///
/// Ok(None) will be returned on EOF, and an error will be returned on read
/// error.
pub fn read_next_rule(input: &mut dyn BufRead) -> Result<Option<String>, std::io::Error> {
    let mut line = String::new();
    loop {
        let mut tmp = String::new();
        let n = input.read_line(&mut tmp)?;
        if n == 0 {
            return Ok(None);
        }

        let tmp = tmp.trim();

        if !tmp.ends_with('\\') {
            line.push_str(tmp);
            break;
        }

        line.push_str(&tmp[..tmp.len() - 1]);
    }
    Ok(Some(line))
}

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

    #[test]
    fn test_parse_token() {
        assert_eq!(parse_header_token("alert"), Ok(("", "alert")));
        assert_eq!(parse_header_token(" alert"), Ok(("", "alert")));
        assert_eq!(parse_header_token(" alert "), Ok((" ", "alert")));
        assert_eq!(parse_header_token("http_uri"), Ok(("", "http_uri")));
    }

    #[test]
    fn test_parse_quoted_string() {
        assert_eq!(
            strip_quotes(r#""some quoted \" string""#),
            r#"some quoted " string"#
        );
    }

    #[test]
    fn test_parse_list() {
        assert_eq!(parse_list_token("[1]"), Ok(("", "[1]")));
        assert_eq!(parse_list_token("[1,2,3]"), Ok(("", "[1,2,3]")));
        assert_eq!(parse_list_token(" [1,2,3]"), Ok(("", "[1,2,3]")));
        assert_eq!(
            parse_list_token(" [1,2,3,[a,b,c]]"),
            Ok(("", "[1,2,3,[a,b,c]]"))
        );

        assert!(parse_list_token("1,2,3]").is_err());
        assert!(parse_list_token("[1,2,3").is_err());

        assert!(parse_list_token("token").is_err());
    }

    #[test]
    fn test_parse_header() {
        let rule = parse_header("alert tcp any any -> any any");
        assert_eq!(
            rule,
            Ok((
                "",
                RuleHeader {
                    action: String::from("alert"),
                    proto: String::from("tcp"),
                    src_addr: String::from("any"),
                    src_port: String::from("any"),
                    direction: String::from("->"),
                    dst_addr: String::from("any"),
                    dst_port: String::from("any"),
                    ..Default::default()
                }
            ))
        );
    }

    #[test]
    fn test_parse_option_without_value() {
        assert_eq!(
            Ok((
                "",
                RuleOption {
                    key: String::from("http_uri"),
                    val: None,
                }
            )),
            parse_option("http_uri;")
        );
    }

    #[test]
    fn test_parse_option() {
        assert_eq!(
            parse_option("msg:value;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("value".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option("msg:value with spaces;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("value with spaces".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option("msg:terminated value with spaces;"),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("terminated value with spaces".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: an escaped \; terminant; next_option;"#),
            Ok((
                " next_option;",
                RuleOption {
                    key: String::from("msg"),
                    val: Some("an escaped ; terminant".to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: "A Quoted Message";"#),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some(r#"A Quoted Message"#.to_string()),
                }
            ))
        );

        assert_eq!(
            parse_option(r#"msg: "A Quoted Message with escaped \" quotes.";"#),
            Ok((
                "",
                RuleOption {
                    key: String::from("msg"),
                    val: Some(r#"A Quoted Message with escaped " quotes."#.to_string()),
                }
            ))
        );
    }

    #[test]
    fn test_parse_rule() {
        let rule = parse_rule(r#"alert ip any any -> any any (msg:"Some Message with a \" Quote"; metadata: key, val; sid:1; rev:1;)"#).unwrap();
        assert_eq!(rule.header.action, "alert");
    }

    #[test]
    fn test_tokenize_option_value() {
        assert_eq!(
            get_option_value("one; two"),
            Ok((" two", String::from("one")))
        );

        assert_eq!(
            get_option_value("one\\;; two"),
            Ok((" two", String::from("one;")))
        );

        let optval = r#""/^(?:[a-zA-Z0-9_%+])*(?:[\x2c\x22\x27\x28]|\x252[c278])/PRi";"#;
        let expected = r#""/^(?:[a-zA-Z0-9_%+])*(?:[\x2c\x22\x27\x28]|\x252[c278])/PRi""#;
        assert_eq!(get_option_value(optval), Ok(("", String::from(expected))));
    }

    #[test]
    fn test_parse_from_reader() {
        let input = r#"alert ip any any -> any any (\
            msg:"TEST RULE"; sid:1; rev:1;)"#;
        let mut reader = input.as_bytes();
        let next = read_next_rule(&mut reader).unwrap();
        assert_eq!(
            next,
            Some(r#"alert ip any any -> any any (msg:"TEST RULE"; sid:1; rev:1;)"#.to_string())
        );
        assert_eq!(read_next_rule(&mut reader).unwrap(), None);
    }
}
<|endoftext|>


<commit_before>use std::fmt::Debug;
use std::fmt;
use crate::model::channel::Message;
use crate::client::Context;
use crate::framework::standard::{Args, CommandOptions};

/// This type describes why a check has failed and occurs on
/// [`CheckResult::Failure`].
///
/// **Note**:
/// The bot-developer is supposed to process this `enum` as the framework is not.
/// It solely serves as a way to inform a user about why a check
/// has failed and for the developer to log given failure (e.g. bugs or statstics)
/// occurring in [`Check`]s.
///
/// [`Check`]: struct.Check.html
/// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
#[derive(Clone, Debug)]
pub enum Reason {
    /// No information on the failure.
    Unknown,
    /// Information dedicated to the user.
    User(String),
    /// Information purely for logging purposes.
    Log(String),
    /// Information for the user but also for logging purposes.
    UserAndLog { user: String, log: String },
    #[doc(hidden)]
    __Nonexhaustive,
}

/// Returned from [`Check`]s.
/// If `Success`, the [`Check`] is considered as passed.
/// If `Failure`, the [`Check`] is considered as failed and can return further
/// information on the cause via [`Reason`].
///
/// [`Check`]: struct.Check.html
/// [`Reason`]: struct.Reason.html
#[derive(Clone, Debug)]
pub enum CheckResult {
   Success,
   Failure(Reason),
}

impl CheckResult {
    /// Creates a new [`CheckResult::Failure`] with [`Reason::User`].
    ///
    /// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
    /// [`Reason::User`]: struct.Reason.html#variant.User
    pub fn new_user<D>(d: D) -> Self
        where D: fmt::Display {
        CheckResult::Failure(Reason::User(d.to_string()))
    }

    /// Creates a new [`CheckResult::Failure`] with [`Reason::Log`].
    ///
    /// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
    /// [`Reason::Log`]: struct.Reason.html#variant.Log
    pub fn new_log<D>(d: D) -> Self
        where D: fmt::Display {
        CheckResult::Failure(Reason::Log(d.to_string()))
    }

    /// Creates a new [`CheckResult::Failure`] with [`Reason::Unknown`].
    ///
    /// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
    /// [`Reason::Unknown`]: struct.Reason.html#variant.Unknown
    pub fn new_unknown() -> Self {
        CheckResult::Failure(Reason::Unknown)
    }

    /// Creates a new [`CheckResult::Failure`] with [`Reason::UserAndLog`].
    ///
    /// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
    /// [`Reason::UserAndLog`]: struct.Reason.html#variant.UserAndLog
    pub fn new_user_and_log<D>(user: D, log: D) -> Self
        where D: fmt::Display {
        CheckResult::Failure(Reason::UserAndLog {
            user: user.to_string(),
            log: log.to_string(),
        })
    }

    /// Returns `true` if [`CheckResult`] is [`CheckResult::Success`] and
    /// `false` if not.
    ///
    /// [`CheckResult`]: enum.CheckResult.html
    /// [`CheckResult::Success`]: enum.CheckResult.html#variant.Success
    pub fn is_success(&self) -> bool {
        if let CheckResult::Success = self {
            return true;
        }

        false
    }
}

impl From<bool> for CheckResult {
    fn from(succeeded: bool) -> Self {
        if succeeded {
            CheckResult::Success
        } else {
            CheckResult::Failure(Reason::Unknown)
        }
    }
}

impl From<Reason> for CheckResult {
    fn from(reason: Reason) -> Self {
        CheckResult::Failure(reason)
    }
}

pub type CheckFunction = fn(&mut Context, &Message, &mut Args, &CommandOptions) -> CheckResult;

/// A check can be part of a command or group and will be executed to
/// determine whether a user is permitted to use related item.
///
/// Additionally, a check may hold additional settings.
pub struct Check {
    /// Name listed in help-system.
    pub name: &'static str,
    /// Function that will be executed.
    pub function: CheckFunction,
    /// Whether a check should be evaluated in the help-system.
    /// `false` will ignore check and won't fail execution.
    pub check_in_help: bool,
    /// Whether a check shall be listed in the help-system.
    /// `false` won't affect whether the check will be evaluated help,
    /// solely `check_in_help` sets this.
    pub display_in_help: bool,
}

impl Debug for Check {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Check")
            .field("name", &self.name)
            .field("function", &"<fn>")
            .field("check_in_help", &self.check_in_help)
            .field("display_in_help", &self.display_in_help)
            .finish()
    }
}

impl PartialEq for Check {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name
    }
}
<commit_msg>Fix `Reason`'s doc-link by using `enum`. (#653)<commit_after>use std::fmt::Debug;
use std::fmt;
use crate::model::channel::Message;
use crate::client::Context;
use crate::framework::standard::{Args, CommandOptions};

/// This type describes why a check has failed and occurs on
/// [`CheckResult::Failure`].
///
/// **Note**:
/// The bot-developer is supposed to process this `enum` as the framework is not.
/// It solely serves as a way to inform a user about why a check
/// has failed and for the developer to log given failure (e.g. bugs or statstics)
/// occurring in [`Check`]s.
///
/// [`Check`]: struct.Check.html
/// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
#[derive(Clone, Debug)]
pub enum Reason {
    /// No information on the failure.
    Unknown,
    /// Information dedicated to the user.
    User(String),
    /// Information purely for logging purposes.
    Log(String),
    /// Information for the user but also for logging purposes.
    UserAndLog { user: String, log: String },
    #[doc(hidden)]
    __Nonexhaustive,
}

/// Returned from [`Check`]s.
/// If `Success`, the [`Check`] is considered as passed.
/// If `Failure`, the [`Check`] is considered as failed and can return further
/// information on the cause via [`Reason`].
///
/// [`Check`]: struct.Check.html
/// [`Reason`]: enum.Reason.html
#[derive(Clone, Debug)]
pub enum CheckResult {
   Success,
   Failure(Reason),
}

impl CheckResult {
    /// Creates a new [`CheckResult::Failure`] with [`Reason::User`].
    ///
    /// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
    /// [`Reason::User`]: enum.Reason.html#variant.User
    pub fn new_user<D>(d: D) -> Self
        where D: fmt::Display {
        CheckResult::Failure(Reason::User(d.to_string()))
    }

    /// Creates a new [`CheckResult::Failure`] with [`Reason::Log`].
    ///
    /// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
    /// [`Reason::Log`]: enum.Reason.html#variant.Log
    pub fn new_log<D>(d: D) -> Self
        where D: fmt::Display {
        CheckResult::Failure(Reason::Log(d.to_string()))
    }

    /// Creates a new [`CheckResult::Failure`] with [`Reason::Unknown`].
    ///
    /// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
    /// [`Reason::Unknown`]: enum.Reason.html#variant.Unknown
    pub fn new_unknown() -> Self {
        CheckResult::Failure(Reason::Unknown)
    }

    /// Creates a new [`CheckResult::Failure`] with [`Reason::UserAndLog`].
    ///
    /// [`CheckResult::Failure`]: enum.CheckResult.html#variant.Failure
    /// [`Reason::UserAndLog`]: enum.Reason.html#variant.UserAndLog
    pub fn new_user_and_log<D>(user: D, log: D) -> Self
        where D: fmt::Display {
        CheckResult::Failure(Reason::UserAndLog {
            user: user.to_string(),
            log: log.to_string(),
        })
    }

    /// Returns `true` if [`CheckResult`] is [`CheckResult::Success`] and
    /// `false` if not.
    ///
    /// [`CheckResult`]: enum.CheckResult.html
    /// [`CheckResult::Success`]: enum.CheckResult.html#variant.Success
    pub fn is_success(&self) -> bool {
        if let CheckResult::Success = self {
            return true;
        }

        false
    }
}

impl From<bool> for CheckResult {
    fn from(succeeded: bool) -> Self {
        if succeeded {
            CheckResult::Success
        } else {
            CheckResult::Failure(Reason::Unknown)
        }
    }
}

impl From<Reason> for CheckResult {
    fn from(reason: Reason) -> Self {
        CheckResult::Failure(reason)
    }
}

pub type CheckFunction = fn(&mut Context, &Message, &mut Args, &CommandOptions) -> CheckResult;

/// A check can be part of a command or group and will be executed to
/// determine whether a user is permitted to use related item.
///
/// Additionally, a check may hold additional settings.
pub struct Check {
    /// Name listed in help-system.
    pub name: &'static str,
    /// Function that will be executed.
    pub function: CheckFunction,
    /// Whether a check should be evaluated in the help-system.
    /// `false` will ignore check and won't fail execution.
    pub check_in_help: bool,
    /// Whether a check shall be listed in the help-system.
    /// `false` won't affect whether the check will be evaluated help,
    /// solely `check_in_help` sets this.
    pub display_in_help: bool,
}

impl Debug for Check {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Check")
            .field("name", &self.name)
            .field("function", &"<fn>")
            .field("check_in_help", &self.check_in_help)
            .field("display_in_help", &self.display_in_help)
            .finish()
    }
}

impl PartialEq for Check {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name
    }
}
<|endoftext|>


<commit_before>//! # `poly_fit` example
//!
//! In this example, we attempt to fit the polynomial f(x) = (x-1) *
//! (x-2) * (x-3) using a set of noisy samples from the interval [0,
//! 4].
//!
//! Specifially, we take 41 equally spaced samples, [0.0, 0.1, 0.2,
//! ..., 3.9, 4.0], We evalutate them on the desired polynomial, and
//! add gaussian noise according to the distrbution N(0, 0.1), to get
//! the RHS matrix `b`.  For the coefficient matrix, we compute [1.0,
//! x, .., x^n] for each sample point `x`, up to some desired degree
//! `n`. As long as `n` >= 3, the resulting solution should be close
//! to the original polynomial.

extern crate rula;
extern crate ndarray;
extern crate rand;

use rula::prelude::*;
use ndarray::prelude::*;
use rand::{thread_rng};
use rand::distributions::{Normal, Sample};

/// Evalutate a polynomial f with coefficients `coefs` at `x`.
fn eval_poly(x: f32, coefs: &[f32]) -> f32 {
    // horner's rule
    coefs.iter().fold(0.0, |acc, c| acc * x + *c)
}

/// Returns the row (1.0, x, x^2, ..., x^n)
fn vandermonde_row(x: f32, n: usize) -> Array<f32, Ix> {
    let mut v: Vec<f32> = Vec::with_capacity(n+1);
    let mut r = 1.0;
    v.push(1.0);
    for i in 1..(n+1) {
        r *= x;
        v.push(r);
    }
    Array::from_vec(v)
}

fn main() {
    let mut rng = thread_rng();
    let mut err = Normal::new(0.0, 0.1);

    // (x-1)*(x-2)*(x-3)
    let coefs = vec![1.0, 11.0, -6.0, -6.0];

    const N: usize = 41;

    let samples = Array::linspace(0.0, 4.0, N);

    const FITTED_POLY_DEGREE: usize = 3;

    // Create the van der monde matrix.
    let mut a = Array::default((N, FITTED_POLY_DEGREE + 1));
    for (x, mut row) in samples.iter().zip(a.outer_iter_mut()) {
        row.assign(&vandermonde_row(*x, FITTED_POLY_DEGREE));
    }

    // Create the 'noisy' solution matrix.
    let mut b = Array::default((N, 1));
    for (i, x) in samples.iter().enumerate() {
        b[(i, 0)] = eval_poly(*x, &coefs) + (err.sample(&mut rng) as f32);
    }

    // Use least squares to fit the matrix.
    let fitted_coefs = LeastSquares::compute_multi(&a, &b);
    assert!(fitted_coefs.is_ok());

    println!("{:?}", fitted_coefs.unwrap().solution);
}
<commit_msg>fix warning in poly_fit example<commit_after>//! # `poly_fit` example
//!
//! In this example, we attempt to fit the polynomial f(x) = (x-1) *
//! (x-2) * (x-3) using a set of noisy samples from the interval [0,
//! 4].
//!
//! Specifially, we take 41 equally spaced samples, [0.0, 0.1, 0.2,
//! ..., 3.9, 4.0], We evalutate them on the desired polynomial, and
//! add gaussian noise according to the distrbution N(0, 0.1), to get
//! the RHS matrix `b`.  For the coefficient matrix, we compute [1.0,
//! x, .., x^n] for each sample point `x`, up to some desired degree
//! `n`. As long as `n` >= 3, the resulting solution should be close
//! to the original polynomial.

extern crate rula;
extern crate ndarray;
extern crate rand;

use rula::prelude::*;
use ndarray::prelude::*;
use rand::{thread_rng};
use rand::distributions::{Normal, Sample};

/// Evalutate a polynomial f with coefficients `coefs` at `x`.
fn eval_poly(x: f32, coefs: &[f32]) -> f32 {
    // horner's rule
    coefs.iter().fold(0.0, |acc, c| acc * x + *c)
}

/// Returns the row (1.0, x, x^2, ..., x^n)
fn vandermonde_row(x: f32, n: usize) -> Array<f32, Ix> {
    let mut v: Vec<f32> = Vec::with_capacity(n+1);
    let mut r = 1.0;
    v.push(1.0);
    for _ in 1..(n+1) {
        r *= x;
        v.push(r);
    }
    Array::from_vec(v)
}

fn main() {
    let mut rng = thread_rng();
    let mut err = Normal::new(0.0, 0.1);

    // (x-1)*(x-2)*(x-3)
    let coefs = vec![1.0, 11.0, -6.0, -6.0];

    const N: usize = 41;

    let samples = Array::linspace(0.0, 4.0, N);

    const FITTED_POLY_DEGREE: usize = 3;

    // Create the van der monde matrix.
    let mut a = Array::default((N, FITTED_POLY_DEGREE + 1));
    for (x, mut row) in samples.iter().zip(a.outer_iter_mut()) {
        row.assign(&vandermonde_row(*x, FITTED_POLY_DEGREE));
    }

    // Create the 'noisy' solution matrix.
    let mut b = Array::default((N, 1));
    for (i, x) in samples.iter().enumerate() {
        b[(i, 0)] = eval_poly(*x, &coefs) + (err.sample(&mut rng) as f32);
    }

    // Use least squares to fit the matrix.
    let fitted_coefs = LeastSquares::compute_multi(&a, &b);
    assert!(fitted_coefs.is_ok());

    println!("{:?}", fitted_coefs.unwrap().solution);
}
<|endoftext|>


<commit_before>use zeroize::Zeroize;

/// Newtype (pattern) to protect secret credentials stored as Strings.
#[derive(Clone)]
pub struct Secret(String);

impl From<String> for Secret {
    fn from(s: String) -> Self {
        Secret(s)
    }
}

/// Allow dereferencing Secrets as &str.
///
/// ```rust
/// # use rusoto_credential::Secret;
/// assert_eq!(Secret::from("hello".to_string()).as_ref(), "hello");
/// ```
impl AsRef<str> for Secret {
    fn as_ref(&self) -> &str {
        return self.0.as_str();
    }
}

impl Drop for Secret {
    fn drop(&mut self) {
        let s = &mut self.0;
        s.zeroize();
    }
}

/// Secrets must not leak, so make sure they are not printed.
///
/// ```rust
/// # use rusoto_credential::Secret;
/// assert_eq!(format!("{:?}",Secret::from("hello".to_string())), "*******");
/// ```
impl std::fmt::Debug for Secret {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "\"*******\"")
    }
}

/// Secrets must not leak, so make sure they are not displayed.
///
/// ```rust
/// # use rusoto_credential::Secret;
/// assert_eq!(format!("{}",Secret::from("hello".to_string())), "*******");
/// ```
impl std::fmt::Display for Secret {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "*******")
    }
}
<commit_msg>Fix doc test for Secret's Debug impl.<commit_after>use zeroize::Zeroize;

/// Newtype (pattern) to protect secret credentials stored as Strings.
#[derive(Clone)]
pub struct Secret(String);

impl From<String> for Secret {
    fn from(s: String) -> Self {
        Secret(s)
    }
}

/// Allow dereferencing Secrets as &str.
///
/// ```rust
/// # use rusoto_credential::Secret;
/// assert_eq!(Secret::from("hello".to_string()).as_ref(), "hello");
/// ```
impl AsRef<str> for Secret {
    fn as_ref(&self) -> &str {
        return self.0.as_str();
    }
}

impl Drop for Secret {
    fn drop(&mut self) {
        let s = &mut self.0;
        s.zeroize();
    }
}

/// Secrets must not leak, so make sure they are not printed.
///
/// ```rust
/// # use rusoto_credential::Secret;
/// assert_eq!(format!("{:?}",Secret::from("hello".to_string())), "\"*******\"");
/// ```
impl std::fmt::Debug for Secret {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "\"*******\"")
    }
}

/// Secrets must not leak, so make sure they are not displayed.
///
/// ```rust
/// # use rusoto_credential::Secret;
/// assert_eq!(format!("{}",Secret::from("hello".to_string())), "*******");
/// ```
impl std::fmt::Display for Secret {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "*******")
    }
}
<|endoftext|>


<commit_before>use scaly::memory::page::_Page;
use scaly::memory::page::_PAGE_SIZE;
use std::alloc::alloc;
use std::alloc::dealloc;
use std::alloc::Layout;
use std::mem::size_of;
use std::ptr::null_mut;

const _STACK_BUCKET_PAGES: usize = 0x100;

pub struct _Region<'a> {
    page: &'a mut _Page,
}

impl<'a> _Region<'a> {
    pub fn create(page: &_Page) -> _Region<'a> {
        let our_page_address = _StackBucket::new_page(page);
        _Region {
            page: unsafe { &mut *(our_page_address as *mut _Page) },
        }
    }
}

impl<'a> Drop for _Region<'a> {
    fn drop(&mut self) {
        self.page.deallocate_extensions();
    }
}

struct _StackBucket {
    next_bucket: *mut _StackBucket,
}

impl _StackBucket {
    pub fn new_page(page: &_Page) -> *mut _Page {
        let page_address = page as *const _Page as usize;
        let stack_bucket_page_address = page_address & !(_PAGE_SIZE * _STACK_BUCKET_PAGES - 1);
        _StackBucket::get_next_page(stack_bucket_page_address, page_address)
    }

    fn get_next_page(stack_bucket_page_address: usize, page_address: usize) -> *mut _Page {
        let overflowed_page_address = stack_bucket_page_address + _PAGE_SIZE * _STACK_BUCKET_PAGES;
        let our_page_address = page_address + _PAGE_SIZE;
        if overflowed_page_address == our_page_address {
            let stack_bucket =
                (stack_bucket_page_address + size_of::<_Page>()) as *mut _StackBucket;
            unsafe { (*stack_bucket).get_page_from_next_bucket() }
        } else {
            our_page_address as *mut _Page
        }
    }

    fn get_page_from_next_bucket(self: &mut _StackBucket) -> *mut _Page {
        if self.next_bucket == null_mut() {
            self.next_bucket = _StackBucket::create();
        }
        _Page::get_page(self.next_bucket as usize)
    }

    fn create() -> *mut _StackBucket {
        unsafe {
            let memory = alloc(Layout::from_size_align_unchecked(
                _PAGE_SIZE * _STACK_BUCKET_PAGES,
                _PAGE_SIZE * _STACK_BUCKET_PAGES,
            ));
            if memory == null_mut() {
                panic!("Unable to crete a new stack bucket: Out of memory.");
            }
            let page = &mut *(memory as *mut _Page);
            page.reset();
            let bucket = page.allocate(_StackBucket {
                next_bucket: null_mut(),
            });
            bucket
        }
    }

    pub fn deallocate(&self) {
        let mut bucket = self as *const _StackBucket;
        unsafe {
            while !(*bucket).next_bucket.is_null() {
                let next_bucket = (*bucket).next_bucket;
                let page = _Page::get_page(bucket as usize);
                dealloc(
                    page as *mut u8,
                    Layout::from_size_align_unchecked(
                        _PAGE_SIZE * _STACK_BUCKET_PAGES,
                        _PAGE_SIZE * _STACK_BUCKET_PAGES,
                    ),
                );

                bucket = next_bucket;
            }
        }
    }
}

impl Drop for _StackBucket {
    fn drop(&mut self) {
        self.deallocate();
    }
}

#[test]
fn test_region() {
    unsafe {
        let root_stack_bucket = _StackBucket::create();
        let r1 = _Region::create(&*_Page::get_page(root_stack_bucket as usize));
        {
            let r2 = _Region::create(r1.page);
            {
                let _ = _Region::create(r2.page);
            }
        }
        (*root_stack_bucket).deallocate();
    }
    unsafe {
        let root_stack_bucket = _StackBucket::create();

        let mut r = _Region::create(&*_Page::get_page(root_stack_bucket as usize));
        for _ in 1.._STACK_BUCKET_PAGES * 2 {
            r = _Region::create(r.page);
        }

        (*root_stack_bucket).deallocate();
    }
}
<commit_msg>Region created from parent region<commit_after>use scaly::memory::page::_Page;
use scaly::memory::page::_PAGE_SIZE;
use std::alloc::alloc;
use std::alloc::dealloc;
use std::alloc::Layout;
use std::mem::size_of;
use std::ptr::null_mut;

const _STACK_BUCKET_PAGES: usize = 0x100;

pub struct _Region<'a> {
    page: &'a mut _Page,
}

impl<'a> _Region<'a> {
    pub fn create_from_page(page: &_Page) -> _Region<'a> {
        let our_page_address = _StackBucket::new_page(page);
        _Region {
            page: unsafe { &mut *(our_page_address as *mut _Page) },
        }
    }
    pub fn create(region: &_Region<'a>) -> _Region<'a> {
        _Region::create_from_page(region.page)
    }
}

impl<'a> Drop for _Region<'a> {
    fn drop(&mut self) {
        self.page.deallocate_extensions();
    }
}

struct _StackBucket {
    next_bucket: *mut _StackBucket,
}

impl _StackBucket {
    pub fn new_page(page: &_Page) -> *mut _Page {
        let page_address = page as *const _Page as usize;
        let stack_bucket_page_address = page_address & !(_PAGE_SIZE * _STACK_BUCKET_PAGES - 1);
        _StackBucket::get_next_page(stack_bucket_page_address, page_address)
    }

    fn get_next_page(stack_bucket_page_address: usize, page_address: usize) -> *mut _Page {
        let overflowed_page_address = stack_bucket_page_address + _PAGE_SIZE * _STACK_BUCKET_PAGES;
        let our_page_address = page_address + _PAGE_SIZE;
        if overflowed_page_address == our_page_address {
            let stack_bucket =
                (stack_bucket_page_address + size_of::<_Page>()) as *mut _StackBucket;
            unsafe { (*stack_bucket).get_page_from_next_bucket() }
        } else {
            our_page_address as *mut _Page
        }
    }

    fn get_page_from_next_bucket(self: &mut _StackBucket) -> *mut _Page {
        if self.next_bucket == null_mut() {
            self.next_bucket = _StackBucket::create();
        }
        _Page::get_page(self.next_bucket as usize)
    }

    fn create() -> *mut _StackBucket {
        unsafe {
            let memory = alloc(Layout::from_size_align_unchecked(
                _PAGE_SIZE * _STACK_BUCKET_PAGES,
                _PAGE_SIZE * _STACK_BUCKET_PAGES,
            ));
            if memory == null_mut() {
                panic!("Unable to crete a new stack bucket: Out of memory.");
            }
            let page = &mut *(memory as *mut _Page);
            page.reset();
            let bucket = page.allocate(_StackBucket {
                next_bucket: null_mut(),
            });
            bucket
        }
    }

    pub fn deallocate(&self) {
        let mut bucket = self as *const _StackBucket;
        unsafe {
            while !(*bucket).next_bucket.is_null() {
                let next_bucket = (*bucket).next_bucket;
                let page = _Page::get_page(bucket as usize);
                dealloc(
                    page as *mut u8,
                    Layout::from_size_align_unchecked(
                        _PAGE_SIZE * _STACK_BUCKET_PAGES,
                        _PAGE_SIZE * _STACK_BUCKET_PAGES,
                    ),
                );

                bucket = next_bucket;
            }
        }
    }
}

impl Drop for _StackBucket {
    fn drop(&mut self) {
        self.deallocate();
    }
}

#[test]
fn test_region() {
    unsafe {
        let root_stack_bucket = _StackBucket::create();
        let r1 = _Region::create_from_page(&*_Page::get_page(root_stack_bucket as usize));
        {
            let r2 = _Region::create(&r1);
            {
                let _ = _Region::create(&r2);
            }
        }
        (*root_stack_bucket).deallocate();
    }
    unsafe {
        let root_stack_bucket = _StackBucket::create();

        let mut r = _Region::create_from_page(&*_Page::get_page(root_stack_bucket as usize));
        for _ in 1.._STACK_BUCKET_PAGES * 2 {
            r = _Region::create(&r);
        }

        (*root_stack_bucket).deallocate();
    }
}
<|endoftext|>