<commit_before>//! Contains various utility types and macros useful for testing hyper clients.
//!
//! # Macros
//! The `mock_connector!` and `mock_connector_in_order!` macros can be used to
//! feed a client with preset data. That way you can control exactly what it will
//! see, confining the test-case to its own sandbox that way.
//!
//! All types they define are public to allow them to be used in other unit-tests.
//! Please note that integration tests cannot share their mock types anyway, as each
//! integration test goes into its own binary.
//!
//! # Usage
//!
//! Set it up for use in tests in `Cargo.toml`
//!
//! ```toml
//! [dev-dependencies]
//! yup-hyper-mock = "*"
//! log = "*"  # log macros are used within yup-hyper-mock
//! ```
//!
//! Link it into your `src/(lib.rs|main.rs)`
//!
//! ```Rust
//! #[cfg(test)] #[macro_use]
//! extern crate "yup-hyper-mock" as hyper_mock
//! ```

extern crate hyper;

#[macro_use]
extern crate log;

use std::fmt;
use std::net::SocketAddr;
use std::io::{self, Read, Write, Cursor};

use hyper::net::{NetworkStream, NetworkConnector};

/// A `NetworkStream` compatible stream that writes into memory, and reads from memory.
pub struct MockStream {
    pub read: Cursor<Vec<u8>>,
    pub write: Vec<u8>,
}

/// A `NetworkStream` compatible stream which contains another `NetworkStream`,
/// whose traffic will be written to another stream.
/// Currently that stream will always be standard error.
pub struct TeeStream<T> {
    pub read_write: T,
    pub copy_to: io::Stderr,
}

impl<T> Clone for TeeStream<T>
    where T: Clone {
    fn clone(&self) -> TeeStream<T> {
        TeeStream {
            read_write: self.read_write.clone(),
            copy_to: io::stderr(),
        }
    }
}

impl<T> Read for TeeStream<T>
    where T: Read {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let res = self.read_write.read(buf);
        match res {
            Ok(s) => {
                self.copy_to.write(&buf[..s]).ok();
            }
            _ => {}
        };
        res
    }
}

impl<T> Write for TeeStream<T>
    where T: Write {
    fn write(&mut self, msg: &[u8]) -> io::Result<usize> {
        self.copy_to.write(msg).ok();
        self.read_write.write(msg)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.copy_to.flush().ok();
        self.read_write.flush()
    }
}

impl<T> NetworkStream for TeeStream<T>
    where T: NetworkStream + Send + Clone {
    fn peer_addr(&mut self) -> io::Result<SocketAddr> {
        self.read_write.peer_addr()
    }
}
impl Clone for MockStream {
    fn clone(&self) -> MockStream {
        MockStream {
            read: Cursor::new(self.read.get_ref().clone()),
            write: self.write.clone()
        }
    }
}

impl fmt::Debug for MockStream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "MockStream {{ read: {:?}, write: {:?} }}", self.read.get_ref(), self.write)
    }
}

impl PartialEq for MockStream {
    fn eq(&self, other: &MockStream) -> bool {
        self.read.get_ref() == other.read.get_ref() && self.write == other.write
    }
}

impl MockStream {
    pub fn new() -> MockStream {
        MockStream {
            read: Cursor::new(vec![]),
            write: vec![],
        }
    }

    pub fn with_input(input: &[u8]) -> MockStream {
        MockStream {
            read: Cursor::new(input.to_vec()),
            write: vec![]
        }
    }
}

impl Read for MockStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.read.read(buf)
    }
}

impl Write for MockStream {
    fn write(&mut self, msg: &[u8]) -> io::Result<usize> {
        Write::write(&mut self.write, msg)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl NetworkStream for MockStream {
    fn peer_addr(&mut self) -> io::Result<SocketAddr> {
        Ok("127.0.0.1:1337".parse().unwrap())
    }
}

/// A `NetworkConnector` which creates `MockStream` instances exclusively.
/// It may be useful to intercept writes.
pub struct MockConnector;

impl NetworkConnector for MockConnector {
    type Stream = MockStream;

    fn connect(&mut self, _host: &str, _port: u16, _scheme: &str) -> io::Result<MockStream> {
        Ok(MockStream::new())
    }
}

/// A `NetworkConnector` embedding another `NetworkConnector` instance,
/// and sets it up to write all reads and writes to standard error as well.
///
/// > **NOTE** It was originally intended to allow arbitrary streams to copy data to,
/// > but I couldn't get passt the compiler with that as normal streams, like files,
/// > are not normally clonable. Maybe an Arc+Mutex would have helped ... .
pub struct TeeConnector<C>
    where C: NetworkConnector {
    pub connector: C,
}

impl<C, S> NetworkConnector for TeeConnector<C>
    where C: NetworkConnector<Stream=S>,
          S: NetworkStream + Send + Clone {
    type Stream = TeeStream<<C as NetworkConnector>::Stream>;

    fn connect(&mut self, _host: &str, _port: u16, _scheme: &str)
        -> io::Result<TeeStream<<C as NetworkConnector>::Stream>> {
        match self.connector.connect(_host, _port, _scheme) {
            Ok(s) => {
                Ok(TeeStream {
                        read_write: s,
                        copy_to: io::stderr(),
                    }
                )
            },
            Err(err) => Err(err),
        }
    }
}

/// This macro maps host URLs to a respective reply, which is given in plain-text.
/// It ignores, but stores, everything that is written to it. However, the stored
/// values are not accessible just yet.
#[macro_export]
macro_rules! mock_connector (
    ($name:ident {
        $($url:expr => $res:expr)*
    }) => (

        pub struct $name;

        impl hyper::net::NetworkConnector for $name {
            type Stream = $crate::MockStream;
            fn connect(&mut self, host: &str, port: u16, scheme: &str) -> ::std::io::Result<$crate::MockStream> {
                use std::collections::HashMap;
                use std::io::Cursor;
                debug!("MockStream::connect({:?}, {:?}, {:?})", host, port, scheme);
                let mut map = HashMap::new();
                $(map.insert($url, $res);)*

                let key = format!("{}://{}", scheme, host);
                // ignore port for now
                match map.get(&*key) {
                    Some(res) => Ok($crate::MockStream {
                        write: vec![],
                        read: Cursor::new(res.to_string().into_bytes()),
                    }),
                    None => panic!("{:?} doesn't know url {}", stringify!($name), key)
                }
            }
        }
    )
);

/// This macro yields all given server replies in the order they are given.
/// The destination host URL doesn't matter at all.
#[macro_export]
macro_rules! mock_connector_in_order (
    ($name:ident {
        $( $res:expr )*
    }) => (

        #[derive(Default)]
        pub struct $name {
            streamers: Vec<String>,
            current: usize,
        }

        impl hyper::net::NetworkConnector for $name {
            type Stream = $crate::MockStream;
            fn connect(&mut self, host: &str, port: u16, scheme: &str) -> ::std::io::Result<$crate::MockStream> {
                use std::io::Cursor;
                debug!("MockStream::connect({:?}, {:?}, {:?})", host, port, scheme);

                if self.streamers.len() == 0 {
                    let mut v = Vec::new();
                    $(v.push($res.to_string());)*
                    self.streamers = v;
                    self.current = 0;
                }
                assert!(self.streamers.len() != 0, "Not a single streamer return value specified");

                let r = Ok($crate::MockStream {
                        write: vec![],
                        read: Cursor::new(self.streamers[self.current].clone().into_bytes())
                });
                self.current += 1;
                r
            }
        }
    )
);

<commit_msg>feat(error): Update to reflect Hyper v0.4<commit_after>//! Contains various utility types and macros useful for testing hyper clients.
//!
//! # Macros
//! The `mock_connector!` and `mock_connector_in_order!` macros can be used to
//! feed a client with preset data. That way you can control exactly what it will
//! see, confining the test-case to its own sandbox that way.
//!
//! All types they define are public to allow them to be used in other unit-tests.
//! Please note that integration tests cannot share their mock types anyway, as each
//! integration test goes into its own binary.
//!
//! # Usage
//!
//! Set it up for use in tests in `Cargo.toml`
//!
//! ```toml
//! [dev-dependencies]
//! yup-hyper-mock = "*"
//! log = "*"  # log macros are used within yup-hyper-mock
//! ```
//!
//! Link it into your `src/(lib.rs|main.rs)`
//!
//! ```Rust
//! #[cfg(test)] #[macro_use]
//! extern crate "yup-hyper-mock" as hyper_mock
//! ```

extern crate hyper;

#[macro_use]
extern crate log;

use std::fmt;
use std::net::SocketAddr;
use std::io::{self, Read, Write, Cursor};

use hyper::net::{NetworkStream, NetworkConnector};

/// A `NetworkStream` compatible stream that writes into memory, and reads from memory.
pub struct MockStream {
    pub read: Cursor<Vec<u8>>,
    pub write: Vec<u8>,
}

/// A `NetworkStream` compatible stream which contains another `NetworkStream`,
/// whose traffic will be written to another stream.
/// Currently that stream will always be standard error.
pub struct TeeStream<T> {
    pub read_write: T,
    pub copy_to: io::Stderr,
}

impl<T> Clone for TeeStream<T>
    where T: Clone {
    fn clone(&self) -> TeeStream<T> {
        TeeStream {
            read_write: self.read_write.clone(),
            copy_to: io::stderr(),
        }
    }
}

impl<T> Read for TeeStream<T>
    where T: Read {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let res = self.read_write.read(buf);
        match res {
            Ok(s) => {
                self.copy_to.write(&buf[..s]).ok();
            }
            _ => {}
        };
        res
    }
}

impl<T> Write for TeeStream<T>
    where T: Write {
    fn write(&mut self, msg: &[u8]) -> io::Result<usize> {
        self.copy_to.write(msg).ok();
        self.read_write.write(msg)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.copy_to.flush().ok();
        self.read_write.flush()
    }
}

impl<T> NetworkStream for TeeStream<T>
    where T: NetworkStream + Send + Clone {
    fn peer_addr(&mut self) -> io::Result<SocketAddr> {
        self.read_write.peer_addr()
    }
}
impl Clone for MockStream {
    fn clone(&self) -> MockStream {
        MockStream {
            read: Cursor::new(self.read.get_ref().clone()),
            write: self.write.clone()
        }
    }
}

impl fmt::Debug for MockStream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "MockStream {{ read: {:?}, write: {:?} }}", self.read.get_ref(), self.write)
    }
}

impl PartialEq for MockStream {
    fn eq(&self, other: &MockStream) -> bool {
        self.read.get_ref() == other.read.get_ref() && self.write == other.write
    }
}

impl MockStream {
    pub fn new() -> MockStream {
        MockStream {
            read: Cursor::new(vec![]),
            write: vec![],
        }
    }

    pub fn with_input(input: &[u8]) -> MockStream {
        MockStream {
            read: Cursor::new(input.to_vec()),
            write: vec![]
        }
    }
}

impl Read for MockStream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.read.read(buf)
    }
}

impl Write for MockStream {
    fn write(&mut self, msg: &[u8]) -> io::Result<usize> {
        Write::write(&mut self.write, msg)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl NetworkStream for MockStream {
    fn peer_addr(&mut self) -> io::Result<SocketAddr> {
        Ok("127.0.0.1:1337".parse().unwrap())
    }
}

/// A `NetworkConnector` which creates `MockStream` instances exclusively.
/// It may be useful to intercept writes.
pub struct MockConnector;

impl NetworkConnector for MockConnector {
    type Stream = MockStream;

    fn connect(&mut self, _host: &str, _port: u16, _scheme: &str) -> hyper::Result<MockStream> {
        Ok(MockStream::new())
    }
}

/// A `NetworkConnector` embedding another `NetworkConnector` instance,
/// and sets it up to write all reads and writes to standard error as well.
///
/// > **NOTE** It was originally intended to allow arbitrary streams to copy data to,
/// > but I couldn't get passt the compiler with that as normal streams, like files,
/// > are not normally clonable. Maybe an Arc+Mutex would have helped ... .
pub struct TeeConnector<C>
    where C: NetworkConnector {
    pub connector: C,
}

impl<C, S> NetworkConnector for TeeConnector<C>
    where C: NetworkConnector<Stream=S>,
          S: NetworkStream + Send + Clone {
    type Stream = TeeStream<<C as NetworkConnector>::Stream>;

    fn connect(&mut self, _host: &str, _port: u16, _scheme: &str)
        -> hyper::Result<TeeStream<<C as NetworkConnector>::Stream>> {
        match self.connector.connect(_host, _port, _scheme) {
            Ok(s) => {
                Ok(TeeStream {
                        read_write: s,
                        copy_to: io::stderr(),
                    }
                )
            },
            Err(err) => Err(err),
        }
    }
}

/// This macro maps host URLs to a respective reply, which is given in plain-text.
/// It ignores, but stores, everything that is written to it. However, the stored
/// values are not accessible just yet.
#[macro_export]
macro_rules! mock_connector (
    ($name:ident {
        $($url:expr => $res:expr)*
    }) => (

        pub struct $name;

        impl hyper::net::NetworkConnector for $name {
            type Stream = $crate::MockStream;
            fn connect(&mut self, host: &str, port: u16, scheme: &str) -> $crate::Result<$crate::MockStream> {
                use std::collections::HashMap;
                use std::io::Cursor;
                debug!("MockStream::connect({:?}, {:?}, {:?})", host, port, scheme);
                let mut map = HashMap::new();
                $(map.insert($url, $res);)*

                let key = format!("{}://{}", scheme, host);
                // ignore port for now
                match map.get(&*key) {
                    Some(res) => Ok($crate::MockStream {
                        write: vec![],
                        read: Cursor::new(res.to_string().into_bytes()),
                    }),
                    None => panic!("{:?} doesn't know url {}", stringify!($name), key)
                }
            }
        }
    )
);

/// This macro yields all given server replies in the order they are given.
/// The destination host URL doesn't matter at all.
#[macro_export]
macro_rules! mock_connector_in_order (
    ($name:ident {
        $( $res:expr )*
    }) => (

        #[derive(Default)]
        pub struct $name {
            streamers: Vec<String>,
            current: usize,
        }

        impl hyper::net::NetworkConnector for $name {
            type Stream = $crate::MockStream;
            fn connect(&mut self, host: &str, port: u16, scheme: &str) -> $crate::Result<$crate::MockStream> {
                use std::io::Cursor;
                debug!("MockStream::connect({:?}, {:?}, {:?})", host, port, scheme);

                if self.streamers.len() == 0 {
                    let mut v = Vec::new();
                    $(v.push($res.to_string());)*
                    self.streamers = v;
                    self.current = 0;
                }
                assert!(self.streamers.len() != 0, "Not a single streamer return value specified");

                let r = Ok($crate::MockStream {
                        write: vec![],
                        read: Cursor::new(self.streamers[self.current].clone().into_bytes())
                });
                self.current += 1;
                r
            }
        }
    )
);

<|endoftext|>


<commit_before>#![feature(asm, alloc_error_handler, core_intrinsics, lang_items, naked_functions)]
#![no_std]

extern crate alloc;

mod callback;

pub mod adc;
pub mod ble_composer;
pub mod ble_parser;
pub mod buttons;
pub mod console;
pub mod debug;
pub mod electronics;
pub mod gpio;
pub mod led;
pub mod result;
pub mod sensors;
pub mod shared_memory;
pub mod simple_ble;
pub mod temperature;
pub mod timer;
pub mod unwind_symbols;

#[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
pub mod entry_point;

#[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
mod lang_items;

#[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
pub mod syscalls;

#[cfg(not(any(target_arch = "arm", target_arch = "riscv32")))]
#[path = "syscalls_mock.rs"]
mod syscalls;

#[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
#[global_allocator]
static ALLOCATOR: linked_list_allocator::LockedHeap = linked_list_allocator::LockedHeap::empty();

// Dummy structure to force importing the panic_handler and other no_std elements when nothing else
// is imported.
pub struct LibTock {}
<commit_msg>Use the no-field syntax instead of empty braces.<commit_after>#![feature(asm, alloc_error_handler, core_intrinsics, lang_items, naked_functions)]
#![no_std]

extern crate alloc;

mod callback;

pub mod adc;
pub mod ble_composer;
pub mod ble_parser;
pub mod buttons;
pub mod console;
pub mod debug;
pub mod electronics;
pub mod gpio;
pub mod led;
pub mod result;
pub mod sensors;
pub mod shared_memory;
pub mod simple_ble;
pub mod temperature;
pub mod timer;
pub mod unwind_symbols;

#[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
pub mod entry_point;

#[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
mod lang_items;

#[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
pub mod syscalls;

#[cfg(not(any(target_arch = "arm", target_arch = "riscv32")))]
#[path = "syscalls_mock.rs"]
mod syscalls;

#[cfg(any(target_arch = "arm", target_arch = "riscv32"))]
#[global_allocator]
static ALLOCATOR: linked_list_allocator::LockedHeap = linked_list_allocator::LockedHeap::empty();

// Dummy structure to force importing the panic_handler and other no_std elements when nothing else
// is imported.
pub struct LibTock;
<|endoftext|>


<commit_before>// error-pattern:cargo-clippy
#![feature(plugin_registrar)]
#![feature(rustc_private)]
#![allow(unknown_lints)]

extern crate rustc_plugin;
use rustc_plugin::Registry;

extern crate clippy_lints;

pub use clippy_lints::*;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    register_plugins(reg);
}

// only exists to let the dogfood integration test works.
// Don't run clippy as an executable directly
#[allow(dead_code, print_stdout)]
fn main() {
    panic!("Please use the cargo-clippy executable");
}
<commit_msg>running cargo clippy on a crate that has the clippy plugin enabled errors out due to duplicate lints<commit_after>// error-pattern:cargo-clippy
#![feature(plugin_registrar)]
#![feature(rustc_private)]
#![allow(unknown_lints)]

extern crate rustc_plugin;
use rustc_plugin::Registry;

extern crate clippy_lints;

pub use clippy_lints::*;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    if reg.sess.lint_store.borrow().get_lint_groups().iter().any(|&(s, _, _)| s == "clippy") {
        reg.sess.struct_warn("running cargo clippy on a crate that also imports the clippy plugin").emit();
    } else {
        register_plugins(reg);
    }
}

// only exists to let the dogfood integration test works.
// Don't run clippy as an executable directly
#[allow(dead_code, print_stdout)]
fn main() {
    panic!("Please use the cargo-clippy executable");
}
<|endoftext|>


<commit_before>#![feature(catch_panic, fnbox)]

//! Work-stealing job-queue implementation.
//! Inspired by the Molecular Matters blog posts.

extern crate crossbeam;
extern crate rand;

use crossbeam::sync::TreiberStack;

use std::any::Any;
use std::boxed::FnBox;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::marker::PhantomData;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Sender, Receiver, TryRecvError};
use std::thread;

use ::rand::thread_rng;
use ::rand::distributions::{IndependentSample, Range};

/// Messages that can be sent to workers.
enum WorkerMessage {
    Stop,
}

type Job = Box<FnBox(&Worker) + 'static + Send>;

/// A job's id in the pool.
pub struct JobId {
    idx: usize,
}

struct PoolEntry {
    job: Job,
    sender: Sender<JobExitStatus>,
}

enum JobExitStatus {
    Success,
    Panic(Box<Any>),
}

/// A mostly lock-free pool allocator of jobs.
/// It only locks when doubling the size of the vector.
/// Eventually, I may implement a lock-free vector based off of
/// Stroustroup-et-al's paper on the subject, but for now, I've
/// decided to "keep it simple, stupid".
struct JobPool {
    jobs: RwLock<Vec<RefCell<Option<PoolEntry>>>>,
    unused: TreiberStack<usize>,
}

impl JobPool {
    /// Initialize this pool with the given capacity.
    fn with_capacity(size: usize) -> Self {
        let mut v = Vec::with_capacity(size);
        let unused = TreiberStack::new();

        for i in (0..size).rev() {
            v.push(RefCell::new(None));

            unused.push(i);
        }

        JobPool {
            jobs: RwLock::new(v),
            unused: unused,
        }
    }

    /// Submit a job to be allocated in the pool.
    fn submit(&self, job: Job) -> (JobId, Receiver<JobExitStatus>) {
        let next = if let Some(next) = self.unused.pop() {
            next
        } else {
            self.double();
            self.unused.pop().unwrap()
        };

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let (tx, rx) = channel();

        let entry = &jobs[next];
        *entry.borrow_mut() = Some(PoolEntry {
            job: job,
            sender: tx,
        });

        (JobId { idx: next, }, rx)
    }

    // double the vector
    fn double(&self) {
        let mut jobs = match self.jobs.write() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let old_len = jobs.len();
        let new_len = old_len * 2;

        jobs.reserve(new_len);

        for i in (old_len..new_len).rev() {
            jobs.push(RefCell::new(None));

            self.unused.push(i);
        }
    }

    /// Get the job given by id. This function should
    /// always yield Some if you use the id you are given.
    fn run(&self, id: JobId, worker: &Worker) {
        use std::mem;

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        // take the job out before pushing the index back onto
        // the unused stack.
        let entry = jobs[id.idx].borrow_mut().take().unwrap();

        let PoolEntry { job, sender } = entry;

        let unbounded_fn: Box<FnBox() + Send> = Box::new(move || job.call_box((worker,)));
        let bounded_fn: Box<FnBox() + 'static + Send> = unsafe { mem::transmute(unbounded_fn) };

        let exit_status = match thread::catch_panic(|| bounded_fn()) {
            Ok(_) => { JobExitStatus::Success }
            Err(e) => { JobExitStatus::Panic(e) }
        };

        // We don't really care if they receive it or not.
        let _result = sender.send(exit_status);
        self.unused.push(id.idx);
    }
}

unsafe impl Send for JobPool {}
unsafe impl Sync for JobPool {}

/// A thread-safe double-ended queue for jobs.
/// Right now, it's just a dumb locking wrapper of VecDeque.
/// It stores indices of jobs in the global job pool.
struct JobQueue {
    inner: Mutex<VecDeque<JobId>>,
}

impl JobQueue {
    fn new() -> Self {
        JobQueue {
            inner: Mutex::new(VecDeque::new()),
        }
    }

    // pop a job from the private end of the queue.
    // only the owner of this queue should call this.
    fn pop(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_back()
    }

    // steal a job from the public end of the queue.
    fn steal(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_front()
    }

    // push another job to the queue.
    fn push(&self, job: JobId) {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.push_back(job);
    }
}

unsafe impl Send for JobQueue {}
unsafe impl Sync for JobQueue {}

/// A worker executes jobs.
pub struct Worker {
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
    siblings: Box<[Arc<JobQueue>]>,
}

impl Worker {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: Some(self),
        }
    }

    fn next_job(&self) -> Option<JobId> {
        if let Some(id) = self.queue.pop() {
            Some(id)
        } else {
            let mut rng = thread_rng();
            let idx = Range::new(0, self.siblings.len()).ind_sample(&mut rng);
            let other_queue: &Arc<JobQueue> = &self.siblings[idx];

            // make sure we don't steal from our own queue.
            if &*self.queue as *const JobQueue != &**other_queue as *const JobQueue {
                other_queue.steal()
            } else {
                None
            }
        }
    }

    fn run_job(&self) {
        if let Some(id) = self.next_job() {
            self.pool.run(id, self);
        }
    }
}

/// A JoinHandle can be used to manually wait for
/// a job to be completed.
pub struct JoinHandle<'a> {
    notifier: Receiver<JobExitStatus>,
    worker: Option<&'a Worker>,
}

impl<'a> JoinHandle<'a> {
    fn wait(self) -> Result<(), Box<Any>> {
        loop {
            match self.notifier.try_recv() {
                Ok(exit_status) => {
                    return match exit_status {
                        JobExitStatus::Success => Ok(()),
                        JobExitStatus::Panic(err) => Err(err),
                    }
                }

                // sender _never_ hangs up before the job is done.
                Err(TryRecvError::Disconnected) => unreachable!(),

                // nothing yet, just do some more work.
                _ => {
                    if let &Some(ref worker) = &self.worker {
                        worker.run_job()
                    }
                }
            }
        }
    }
}

pub struct JobSystem {
    handles: Vec<Sender<WorkerMessage>>,
    worker: Worker,
}

impl JobSystem {
    pub fn get_worker(&self) -> &Worker {
        &self.worker
    }
}

impl Drop for JobSystem {
    fn drop(&mut self) {
        for handle in &self.handles {
            match handle.send(WorkerMessage::Stop) {
                _ => {}
            }
        }
    }
}

/// Create a new `JobSystem` with `n_threads` workers, including the
/// current thread. This will fail if `n_threads` is 0.
pub fn make_pool(n_threads: usize) -> Option<JobSystem> {
    if n_threads == 0 {
        return None;
    }

    let job_pool = Arc::new(JobPool::with_capacity(4096));

    // make a queue for every thread, including this one.
    let queues: Vec<_> = (0..n_threads)
        .map(|_| Arc::new(JobQueue::new())).collect();

    let queues = queues.into_boxed_slice();

    let mut handles = Vec::new();

    // Spawn a worker thread for every queue but the last one,
    // since that belongs to the thread this is spawned on.
    for queue in queues.iter().cloned().take(queues.len() - 1) {
        let (tx, rx) = channel();
        let siblings = queues.clone();
        let pool = job_pool.clone();

        thread::spawn(move || {
            let worker = Worker {
                pool: pool,
                queue: queue,
                siblings: siblings,
            };

            worker_main(worker, rx);
        });
        handles.push(tx);
    }

    Some(JobSystem {
        handles: handles,
        worker: Worker {
            pool: job_pool,
            queue: queues.last().unwrap().clone(),
            siblings: queues,
        },
    })
}

// the main thread for each worker
fn worker_main(worker: Worker, rx: Receiver<WorkerMessage>) {
    loop {
        // check for message
        match rx.try_recv() {
            // time to stop!
            Ok(WorkerMessage::Stop) => {
                break;
            }

            // somehow the JobSystem got dropped without disconnecting.
            Err(TryRecvError::Disconnected) => {
                break;
            }

            _ => {}
        }

        // run a job
        worker.run_job();
    }
}

#[cfg(test)]
mod tests {
    use std::thread;

    use super::*;

    #[test]
    fn basic_spawning() {
        // four worker threads, including this one.
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();

        let mut my_jobs = Vec::new();

        for _ in 0..100 {
            let handle = this_worker.submit(|_| println!("Hello!"));
            my_jobs.push(handle);
        }
    }

    #[test]
    fn child_jobs() {
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();

        this_worker.submit(|worker| {
            let mut jobs = Vec::new();
            for _ in 0..100 {
                let handle = worker.submit(|_| println!("Hello!"));
                jobs.push(handle);
            };
        }).wait();
    }

/*    #[test]
    fn scoping() {
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();
        this_worker.submit(|worker| {
            let mut v = vec![0; 256];

            worker.scope(|scope| {
                for chunk in v.chunks_mut(32) {
                    scope.submit(move |_| {
                        for i in chunk { *i += 1 }
                    });
                }
            });

            for i in v {
                assert_eq!(i, 1);
            }
        }).wait();
    }*/
}<commit_msg>Smashed warnings<commit_after>#![feature(catch_panic, fnbox)]

//! Work-stealing job-queue implementation.
//! Inspired by the Molecular Matters blog posts.

extern crate crossbeam;
extern crate rand;

use crossbeam::sync::TreiberStack;

use std::any::Any;
use std::boxed::FnBox;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver, TryRecvError};
use std::thread;

use ::rand::thread_rng;
use ::rand::distributions::{IndependentSample, Range};

/// Messages that can be sent to workers.
enum WorkerMessage {
    Stop,
}

type Job = Box<FnBox(&Worker) + 'static + Send>;

/// A job's id in the pool.
pub struct JobId {
    idx: usize,
}

struct PoolEntry {
    job: Job,
    sender: Sender<JobExitStatus>,
}

enum JobExitStatus {
    Success,
    Panic(Box<Any>),
}

/// A mostly lock-free pool allocator of jobs.
/// It only locks when doubling the size of the vector.
/// Eventually, I may implement a lock-free vector based off of
/// Stroustroup-et-al's paper on the subject, but for now, I've
/// decided to "keep it simple, stupid".
struct JobPool {
    jobs: RwLock<Vec<RefCell<Option<PoolEntry>>>>,
    unused: TreiberStack<usize>,
}

impl JobPool {
    /// Initialize this pool with the given capacity.
    fn with_capacity(size: usize) -> Self {
        let mut v = Vec::with_capacity(size);
        let unused = TreiberStack::new();

        for i in (0..size).rev() {
            v.push(RefCell::new(None));

            unused.push(i);
        }

        JobPool {
            jobs: RwLock::new(v),
            unused: unused,
        }
    }

    /// Submit a job to be allocated in the pool.
    fn submit(&self, job: Job) -> (JobId, Receiver<JobExitStatus>) {
        let next = if let Some(next) = self.unused.pop() {
            next
        } else {
            self.double();
            self.unused.pop().unwrap()
        };

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let (tx, rx) = channel();

        let entry = &jobs[next];
        *entry.borrow_mut() = Some(PoolEntry {
            job: job,
            sender: tx,
        });

        (JobId { idx: next, }, rx)
    }

    // double the vector
    fn double(&self) {
        let mut jobs = match self.jobs.write() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        let old_len = jobs.len();
        let new_len = old_len * 2;

        jobs.reserve(new_len);

        for i in (old_len..new_len).rev() {
            jobs.push(RefCell::new(None));

            self.unused.push(i);
        }
    }

    /// Get the job given by id. This function should
    /// always yield Some if you use the id you are given.
    fn run(&self, id: JobId, worker: &Worker) {
        use std::mem;

        let jobs = match self.jobs.read() {
            Ok(j) => j,
            Err(p) => p.into_inner(),
        };

        // take the job out before pushing the index back onto
        // the unused stack.
        let entry = jobs[id.idx].borrow_mut().take().unwrap();

        let PoolEntry { job, sender } = entry;

        let unbounded_fn: Box<FnBox() + Send> = Box::new(move || job.call_box((worker,)));
        let bounded_fn: Box<FnBox() + 'static + Send> = unsafe { mem::transmute(unbounded_fn) };

        let exit_status = match thread::catch_panic(|| bounded_fn()) {
            Ok(_) => { JobExitStatus::Success }
            Err(e) => { JobExitStatus::Panic(e) }
        };

        // We don't really care if they receive it or not.
        let _result = sender.send(exit_status);
        self.unused.push(id.idx);
    }
}

unsafe impl Send for JobPool {}
unsafe impl Sync for JobPool {}

/// A thread-safe double-ended queue for jobs.
/// Right now, it's just a dumb locking wrapper of VecDeque.
/// It stores indices of jobs in the global job pool.
struct JobQueue {
    inner: Mutex<VecDeque<JobId>>,
}

impl JobQueue {
    fn new() -> Self {
        JobQueue {
            inner: Mutex::new(VecDeque::new()),
        }
    }

    // pop a job from the private end of the queue.
    // only the owner of this queue should call this.
    fn pop(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_back()
    }

    // steal a job from the public end of the queue.
    fn steal(&self) -> Option<JobId> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.pop_front()
    }

    // push another job to the queue.
    fn push(&self, job: JobId) {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(p) => p.into_inner(),
        };

        guard.push_back(job);
    }
}

unsafe impl Send for JobQueue {}
unsafe impl Sync for JobQueue {}

/// A worker executes jobs.
pub struct Worker {
    pool: Arc<JobPool>,
    queue: Arc<JobQueue>,
    siblings: Box<[Arc<JobQueue>]>,
}

impl Worker {
    pub fn submit<F: FnOnce(&Worker) + 'static + Send>(&self, f: F) -> JoinHandle {
        let job = Box::new(f);
        let (id, notifier) = self.pool.submit(job);
        self.queue.push(id);

        JoinHandle {
            notifier: notifier,
            worker: Some(self),
        }
    }

    fn next_job(&self) -> Option<JobId> {
        if let Some(id) = self.queue.pop() {
            Some(id)
        } else {
            let mut rng = thread_rng();
            let idx = Range::new(0, self.siblings.len()).ind_sample(&mut rng);
            let other_queue: &Arc<JobQueue> = &self.siblings[idx];

            // make sure we don't steal from our own queue.
            if &*self.queue as *const JobQueue != &**other_queue as *const JobQueue {
                other_queue.steal()
            } else {
                None
            }
        }
    }

    fn run_job(&self) {
        if let Some(id) = self.next_job() {
            self.pool.run(id, self);
        }
    }
}

/// A JoinHandle can be used to manually wait for
/// a job to be completed.
pub struct JoinHandle<'a> {
    notifier: Receiver<JobExitStatus>,
    worker: Option<&'a Worker>,
}

impl<'a> JoinHandle<'a> {
    pub fn wait(self) -> Result<(), Box<Any>> {
        loop {
            match self.notifier.try_recv() {
                Ok(exit_status) => {
                    return match exit_status {
                        JobExitStatus::Success => Ok(()),
                        JobExitStatus::Panic(err) => Err(err),
                    }
                }

                // sender _never_ hangs up before the job is done.
                Err(TryRecvError::Disconnected) => unreachable!(),

                // nothing yet, just do some more work.
                _ => {
                    if let &Some(ref worker) = &self.worker {
                        worker.run_job()
                    }
                }
            }
        }
    }
}

pub struct JobSystem {
    handles: Vec<Sender<WorkerMessage>>,
    worker: Worker,
}

impl JobSystem {
    pub fn get_worker(&self) -> &Worker {
        &self.worker
    }
}

impl Drop for JobSystem {
    fn drop(&mut self) {
        for handle in &self.handles {
            match handle.send(WorkerMessage::Stop) {
                _ => {}
            }
        }
    }
}

/// Create a new `JobSystem` with `n_threads` workers, including the
/// current thread. This will fail if `n_threads` is 0.
pub fn make_pool(n_threads: usize) -> Option<JobSystem> {
    if n_threads == 0 {
        return None;
    }

    let job_pool = Arc::new(JobPool::with_capacity(4096));

    // make a queue for every thread, including this one.
    let queues: Vec<_> = (0..n_threads)
        .map(|_| Arc::new(JobQueue::new())).collect();

    let queues = queues.into_boxed_slice();

    let mut handles = Vec::new();

    // Spawn a worker thread for every queue but the last one,
    // since that belongs to the thread this is spawned on.
    for queue in queues.iter().cloned().take(queues.len() - 1) {
        let (tx, rx) = channel();
        let siblings = queues.clone();
        let pool = job_pool.clone();

        thread::spawn(move || {
            let worker = Worker {
                pool: pool,
                queue: queue,
                siblings: siblings,
            };

            worker_main(worker, rx);
        });
        handles.push(tx);
    }

    Some(JobSystem {
        handles: handles,
        worker: Worker {
            pool: job_pool,
            queue: queues.last().unwrap().clone(),
            siblings: queues,
        },
    })
}

// the main thread for each worker
fn worker_main(worker: Worker, rx: Receiver<WorkerMessage>) {
    loop {
        // check for message
        match rx.try_recv() {
            // time to stop!
            Ok(WorkerMessage::Stop) => {
                break;
            }

            // somehow the JobSystem got dropped without disconnecting.
            Err(TryRecvError::Disconnected) => {
                break;
            }

            _ => {}
        }

        // run a job
        worker.run_job();
    }
}

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

    #[test]
    fn basic_spawning() {
        // four worker threads, including this one.
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();

        let mut my_jobs = Vec::new();

        for _ in 0..100 {
            let handle = this_worker.submit(|_| println!("Hello!"));
            my_jobs.push(handle);
        }
    }

    #[test]
    fn child_jobs() {
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();

        this_worker.submit(|worker| {
            let mut jobs = Vec::new();
            for _ in 0..100 {
                let handle = worker.submit(|_| println!("Hello!"));
                jobs.push(handle);
            };
        }).wait().unwrap();
    }

/*    #[test]
    fn scoping() {
        let pool = make_pool(4).unwrap();
        let this_worker = pool.get_worker();
        this_worker.submit(|worker| {
            let mut v = vec![0; 256];

            worker.scope(|scope| {
                for chunk in v.chunks_mut(32) {
                    scope.submit(move |_| {
                        for i in chunk { *i += 1 }
                    });
                }
            });

            for i in v {
                assert_eq!(i, 1);
            }
        }).wait();
    }*/
}<|endoftext|>


<commit_before>// Copyright 2015 Keegan McAllister.
// Copyright 2016 Corey Farwell.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// See `LICENSE` in this repository.

use std::io::{self, Read};
use std::{panic, process};

mod common;
pub use common::*;

/// Utility that reads a `Vec` of bytes from standard input (stdin)
/// and passes it to `closure`. All panics that occur within
/// `closure` will be treated as aborts. This is done so that
/// AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Image;
/// # impl Image {
/// #     fn parse(_: &[u8]) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_bytes(|read| {
///         Image::parse(&read)
///     })
/// }
/// ```
#[deprecated(
    since = "0.3.3",
    note = "This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead."
)]
pub fn read_stdio_bytes<F>(closure: F)
where
    F: Fn(Vec<u8>) + panic::RefUnwindSafe,
{
    let mut input = vec![];
    let result = io::stdin().read_to_end(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

/// Utility that reads a `String` from standard input (stdin) and
/// passes it to `closure`. If a `String` cannot be constructed from
/// the data provided by standard input, `closure` will _not_ be
/// called. All panics that occur within `closure` will be treated as
/// aborts. This is done so that AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Url;
/// # impl Url {
/// #     fn parse(_: &str) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_string(|string| {
///         Url::parse(&string)
///     })
/// }
/// ```
#[deprecated(
    since = "0.3.3",
    note = "This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead."
)]
pub fn read_stdio_string<F>(closure: F)
where
    F: Fn(String) + panic::RefUnwindSafe,
{
    let mut input = String::new();
    let result = io::stdin().read_to_string(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

// those functions are provided by the afl-llvm-rt static library
extern "C" {
    fn __afl_persistent_loop(counter: usize) -> isize;
    fn __afl_manual_init();
}

/// Fuzz a closure by passing it a `&[u8]`
///
/// This slice contains a "random" quantity of "random" data.
///
/// ```rust,no_run
/// # extern crate afl;
/// # use afl::fuzz;
/// # fn main() {
/// fuzz(true, |data|{
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
pub fn fuzz<F>(hook: bool, mut closure: F)
where
    F: FnMut(&[u8]) + std::panic::RefUnwindSafe,
{
    // this marker strings needs to be in the produced executable for
    // afl-fuzz to detect `persistent mode` and `defered mode`
    static PERSIST_MARKER: &'static str = "##SIG_AFL_PERSISTENT##\0";
    static DEFERED_MARKER: &'static str = "##SIG_AFL_DEFER_FORKSRV##\0";

    // we now need a fake instruction to prevent the compiler from optimizing out
    // those marker strings
    unsafe { std::ptr::read_volatile(&PERSIST_MARKER) }; // hack used in https://github.com/bluss/bencher for black_box()
    unsafe { std::ptr::read_volatile(&DEFERED_MARKER) };
    // unsafe { asm!("" : : "r"(&PERSIST_MARKER)) }; // hack used in nightly's back_box(), requires feature asm
    // unsafe { asm!("" : : "r"(&DEFERED_MARKER)) };

    if hook {
        // sets panic hook to abort
        std::panic::set_hook(Box::new(|_| {
            std::process::abort();
        }));
    }

    let mut input = vec![];

    // initialize forkserver there
    unsafe { __afl_manual_init() };

    while unsafe { __afl_persistent_loop(1000) } != 0 {
        // get buffer from AFL through stdin
        let result = io::stdin().read_to_end(&mut input);
        if result.is_err() {
            return;
        }

        // We still catch unwinding panics just in case the fuzzed code modifies
        // the panic hook.
        // If so, the fuzzer will be unable to tell different bugs appart and you will
        // only be able to find one bug at a time before fixing it to then find a new one.
        let did_panic = std::panic::catch_unwind(panic::AssertUnwindSafe(|| {
            closure(&input);
        }))
        .is_err();

        if did_panic {
            // hopefully the custom panic hook will be called before and abort the
            // process before the stack frames are unwinded.
            std::process::abort();
        }
        // The version of lazy_static.rs at https://github.com/rust-fuzz/resettable-lazy-static.rs
        // provides the `reset` function.
        #[cfg(feature = "reset_lazy_static")]
        unsafe {
            lazy_static::lazy::reset();
        }
        input.clear();
    }
}

/// Fuzz a closure-like block of code by passing it an object of arbitrary type.
///
/// You can choose the type of the argument using the syntax as in the example below.
/// Please check out the `arbitrary` crate to see which types are available.
///
/// For performance reasons, it is recommended that you use the native type `&[u8]` when possible.
///
/// ```rust,no_run
/// # #[macro_use] extern crate afl;
/// # fn main() {
/// fuzz!(|data: &[u8]| {
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
#[macro_export]
macro_rules! fuzz {
    ( $($x:tt)* ) => { afl::__fuzz!(true, $($x)*) }
}

/// Like `fuzz!` above, but panics that are caught inside the fuzzed code are not turned into
/// crashes.
#[macro_export]
macro_rules! fuzz_nohook {
    ( $($x:tt)* ) => { afl::__fuzz!(false, $($x)*) }
}

#[macro_export]
macro_rules! __fuzz {
    ($hook:expr, |$buf:ident| $body:block) => {
        afl::fuzz($hook, |$buf| $body);
    };
    ($hook:expr, |$buf:ident: &[u8]| $body:block) => {
        afl::fuzz($hook, |$buf| $body);
    };
    ($hook:expr, |$buf:ident: $dty: ty| $body:block) => {
        afl::fuzz($hook, |$buf| {
            let $buf: $dty = {
                use arbitrary::{Arbitrary, RingBuffer};
                if let Ok(d) = RingBuffer::new($buf, $buf.len())
                    .and_then(|mut b| Arbitrary::arbitrary(&mut b).map_err(|_| ""))
                {
                    d
                } else {
                    return;
                }
            };

            $body
        });
    };
}

#[cfg(test)]
mod test {
    /*
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::thread;
    use std::time;

    extern crate libc;
    extern crate tempdir;

    fn target_path() -> PathBuf {
        if PathBuf::from("../target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("../target/debug/")
        } else if PathBuf::from("target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("target/debug/")
        } else {
            panic!("Could not find cargo-afl-fuzz!");
        }
    }

    #[test]
    fn test_cargo_afl_fuzz() {
        let temp_dir = tempdir::TempDir::new("aflrs").expect("Could not create temporary directory");
        let temp_dir_path = temp_dir.path();
        let mut child = Command::new(target_path().join("cargo-afl-fuzz"))
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .arg("-i")
            .arg(".")
            .arg("-o")
            .arg(temp_dir_path)
            .arg(target_path().join("examples").join("hello"))
            .spawn()
            .expect("Could not run cargo-afl-fuzz");
        thread::sleep(time::Duration::from_secs(7));
        child.kill().unwrap();
        assert!(temp_dir_path.join("fuzzer_stats").is_file());
    }
    */
}
<commit_msg>Use `$crate::` instead of `afl::`<commit_after>// Copyright 2015 Keegan McAllister.
// Copyright 2016 Corey Farwell.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// See `LICENSE` in this repository.

use std::io::{self, Read};
use std::{panic, process};

mod common;
pub use common::*;

/// Utility that reads a `Vec` of bytes from standard input (stdin)
/// and passes it to `closure`. All panics that occur within
/// `closure` will be treated as aborts. This is done so that
/// AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Image;
/// # impl Image {
/// #     fn parse(_: &[u8]) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_bytes(|read| {
///         Image::parse(&read)
///     })
/// }
/// ```
#[deprecated(
    since = "0.3.3",
    note = "This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead."
)]
pub fn read_stdio_bytes<F>(closure: F)
where
    F: Fn(Vec<u8>) + panic::RefUnwindSafe,
{
    let mut input = vec![];
    let result = io::stdin().read_to_end(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

/// Utility that reads a `String` from standard input (stdin) and
/// passes it to `closure`. If a `String` cannot be constructed from
/// the data provided by standard input, `closure` will _not_ be
/// called. All panics that occur within `closure` will be treated as
/// aborts. This is done so that AFL considers a panic to be a crash.
///
/// # Examples
///
/// ```no_run
/// extern crate afl;
/// # struct Url;
/// # impl Url {
/// #     fn parse(_: &str) {}
/// # }
///
/// fn main() {
///     afl::read_stdio_string(|string| {
///         Url::parse(&string)
///     })
/// }
/// ```
#[deprecated(
    since = "0.3.3",
    note = "This function does not use the `persistent mode` and `defered forkserver mode` and is therefore very slow.
Please use fuzz() or fuzz!() instead."
)]
pub fn read_stdio_string<F>(closure: F)
where
    F: Fn(String) + panic::RefUnwindSafe,
{
    let mut input = String::new();
    let result = io::stdin().read_to_string(&mut input);
    if result.is_err() {
        return;
    }
    let was_panic = panic::catch_unwind(|| {
        closure(input);
    });
    if was_panic.is_err() {
        process::abort();
    }
}

// those functions are provided by the afl-llvm-rt static library
extern "C" {
    fn __afl_persistent_loop(counter: usize) -> isize;
    fn __afl_manual_init();
}

/// Fuzz a closure by passing it a `&[u8]`
///
/// This slice contains a "random" quantity of "random" data.
///
/// ```rust,no_run
/// # extern crate afl;
/// # use afl::fuzz;
/// # fn main() {
/// fuzz(true, |data|{
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
pub fn fuzz<F>(hook: bool, mut closure: F)
where
    F: FnMut(&[u8]) + std::panic::RefUnwindSafe,
{
    // this marker strings needs to be in the produced executable for
    // afl-fuzz to detect `persistent mode` and `defered mode`
    static PERSIST_MARKER: &'static str = "##SIG_AFL_PERSISTENT##\0";
    static DEFERED_MARKER: &'static str = "##SIG_AFL_DEFER_FORKSRV##\0";

    // we now need a fake instruction to prevent the compiler from optimizing out
    // those marker strings
    unsafe { std::ptr::read_volatile(&PERSIST_MARKER) }; // hack used in https://github.com/bluss/bencher for black_box()
    unsafe { std::ptr::read_volatile(&DEFERED_MARKER) };
    // unsafe { asm!("" : : "r"(&PERSIST_MARKER)) }; // hack used in nightly's back_box(), requires feature asm
    // unsafe { asm!("" : : "r"(&DEFERED_MARKER)) };

    if hook {
        // sets panic hook to abort
        std::panic::set_hook(Box::new(|_| {
            std::process::abort();
        }));
    }

    let mut input = vec![];

    // initialize forkserver there
    unsafe { __afl_manual_init() };

    while unsafe { __afl_persistent_loop(1000) } != 0 {
        // get buffer from AFL through stdin
        let result = io::stdin().read_to_end(&mut input);
        if result.is_err() {
            return;
        }

        // We still catch unwinding panics just in case the fuzzed code modifies
        // the panic hook.
        // If so, the fuzzer will be unable to tell different bugs appart and you will
        // only be able to find one bug at a time before fixing it to then find a new one.
        let did_panic = std::panic::catch_unwind(panic::AssertUnwindSafe(|| {
            closure(&input);
        }))
        .is_err();

        if did_panic {
            // hopefully the custom panic hook will be called before and abort the
            // process before the stack frames are unwinded.
            std::process::abort();
        }
        // The version of lazy_static.rs at https://github.com/rust-fuzz/resettable-lazy-static.rs
        // provides the `reset` function.
        #[cfg(feature = "reset_lazy_static")]
        unsafe {
            lazy_static::lazy::reset();
        }
        input.clear();
    }
}

/// Fuzz a closure-like block of code by passing it an object of arbitrary type.
///
/// You can choose the type of the argument using the syntax as in the example below.
/// Please check out the `arbitrary` crate to see which types are available.
///
/// For performance reasons, it is recommended that you use the native type `&[u8]` when possible.
///
/// ```rust,no_run
/// # #[macro_use] extern crate afl;
/// # fn main() {
/// fuzz!(|data: &[u8]| {
///     if data.len() != 6 {return}
///     if data[0] != b'q' {return}
///     if data[1] != b'w' {return}
///     if data[2] != b'e' {return}
///     if data[3] != b'r' {return}
///     if data[4] != b't' {return}
///     if data[5] != b'y' {return}
///     panic!("BOOM")
/// });
/// # }
/// ```
#[macro_export]
macro_rules! fuzz {
    ( $($x:tt)* ) => { $crate::__fuzz!(true, $($x)*) }
}

/// Like `fuzz!` above, but panics that are caught inside the fuzzed code are not turned into
/// crashes.
#[macro_export]
macro_rules! fuzz_nohook {
    ( $($x:tt)* ) => { $crate::__fuzz!(false, $($x)*) }
}

#[macro_export]
macro_rules! __fuzz {
    ($hook:expr, |$buf:ident| $body:block) => {
        $crate::fuzz($hook, |$buf| $body);
    };
    ($hook:expr, |$buf:ident: &[u8]| $body:block) => {
        $crate::fuzz($hook, |$buf| $body);
    };
    ($hook:expr, |$buf:ident: $dty: ty| $body:block) => {
        $crate::fuzz($hook, |$buf| {
            let $buf: $dty = {
                use arbitrary::{Arbitrary, RingBuffer};
                if let Ok(d) = RingBuffer::new($buf, $buf.len())
                    .and_then(|mut b| Arbitrary::arbitrary(&mut b).map_err(|_| ""))
                {
                    d
                } else {
                    return;
                }
            };

            $body
        });
    };
}

#[cfg(test)]
mod test {
    /*
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::thread;
    use std::time;

    extern crate libc;
    extern crate tempdir;

    fn target_path() -> PathBuf {
        if PathBuf::from("../target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("../target/debug/")
        } else if PathBuf::from("target/debug/cargo-afl-fuzz").exists() {
            PathBuf::from("target/debug/")
        } else {
            panic!("Could not find cargo-afl-fuzz!");
        }
    }

    #[test]
    fn test_cargo_afl_fuzz() {
        let temp_dir = tempdir::TempDir::new("aflrs").expect("Could not create temporary directory");
        let temp_dir_path = temp_dir.path();
        let mut child = Command::new(target_path().join("cargo-afl-fuzz"))
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .arg("-i")
            .arg(".")
            .arg("-o")
            .arg(temp_dir_path)
            .arg(target_path().join("examples").join("hello"))
            .spawn()
            .expect("Could not run cargo-afl-fuzz");
        thread::sleep(time::Duration::from_secs(7));
        child.kill().unwrap();
        assert!(temp_dir_path.join("fuzzer_stats").is_file());
    }
    */
}
<|endoftext|>


<commit_before>extern crate byteorder;
#[macro_use(_tt_as_expr_hack)]
#[macro_use(matches)] extern crate matches;
#[macro_use(quick_error)] extern crate quick_error;

mod enums;
mod structs;
mod name;
mod parser;
mod error;
mod header;
mod rrdata;
mod builder;

pub use enums::{Type, QueryType, Class, QueryClass, ResponseCode, Opcode};
pub use structs::{Question, ResourceRecord, OptRecord, Packet, SoaRecord};
pub use name::{Name};
pub use error::{Error};
pub use header::{Header};
pub use rrdata::{RRData};
pub use builder::{Builder};
<commit_msg>Import all macros from `matches`<commit_after>extern crate byteorder;
#[macro_use] extern crate matches;
#[macro_use(quick_error)] extern crate quick_error;

mod enums;
mod structs;
mod name;
mod parser;
mod error;
mod header;
mod rrdata;
mod builder;

pub use enums::{Type, QueryType, Class, QueryClass, ResponseCode, Opcode};
pub use structs::{Question, ResourceRecord, OptRecord, Packet, SoaRecord};
pub use name::{Name};
pub use error::{Error};
pub use header::{Header};
pub use rrdata::{RRData};
pub use builder::{Builder};
<|endoftext|>


<commit_before>extern crate toml;
extern crate shellexpand;

// Import traits
use std::io::prelude::Read;

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}

pub fn test() {
    println!("Hello from lib.rs");
}

#[derive(Debug)]
enum OnSuccess {
    Continue,
    Echo {message: String},
    Warn {message: String},
    Abort {message: String},
}

#[derive(Debug)]
enum OnFailure {
    Continue,
    Echo {message: String},
    Warn {message: String},
    Abort {message: String},
}

#[derive(Debug)]
enum Pipe {
    Null,
    StdOut,
    StdErr,
    Variable {name: String},
    File {filename: String},
}

#[derive(Debug)]
enum Step {
    Echo {name: String, message: String},
    Run {name: String, command: String, arguments: String, on_success: OnSuccess, on_failure: OnFailure},
    Shell {name: String, command: String, stdout: Pipe, on_success: OnSuccess, on_failure: OnFailure},
    SetEnv {variable: String, value: String},
}

fn parse_toml(filename: &str) -> Vec<Step> {

    println!("Reading toml file {}", filename);

    // Read toml file content
    let mut input = String::new();
    std::fs::File::open(&filename).and_then(|mut f| {
        f.read_to_string(&mut input)
    }).unwrap();

    // Pase toml file
    let mut parser = toml::Parser::new(&input);

    let toml = parser.parse().unwrap_or_else(|| {
        for err in &parser.errors {
            let (loline, locol) = parser.to_linecol(err.lo);
            let (hiline, hicol) = parser.to_linecol(err.hi);
            println!("{}:{}:{}-{}:{} error: {}",
                     filename, loline, locol, hiline, hicol, err.desc);
        }
        panic!("Error parsing toml file.");
    });

    // println!("toml: {:#?}", toml);

    // Parse "variables" table
    for (variable, value) in toml.get("variables").unwrap().as_table().unwrap() {
        let value = value.as_str().unwrap();
        // println!("variable: {}   value: {}", variable, value);
        std::env::set_var(variable, value);
    }

    // Parse "steps" table
    let mut steps = Vec::<Step>::with_capacity(10);

    for step in toml.get("step").unwrap().as_slice().unwrap() {
        let action = step.as_table().unwrap().get("action").unwrap().as_str().unwrap();
        let name = step.as_table().unwrap().get("name").unwrap().as_str().unwrap().to_string();
        let details = step.as_table().unwrap().get("details").unwrap().as_table().unwrap();
        let step = match action {
            "echo" => {
                let message = details.get("message").unwrap().as_str().unwrap().to_string();
                Step::Echo {name: name, message: message}
            },
            "run" => {
                let command = details.get("command").unwrap().as_str().unwrap().to_string();

                // Allow non-existing 'arguments' field, defaulting to empty string.
                let arguments_option = details.get("arguments");
                let arguments = match arguments_option {
                    None => String::new(),
                    Some(arguments_toml_value) => arguments_toml_value.as_str().unwrap().to_string(),
                };

                let on_success = match details.get("on_success") {
                    None => {
                        println!("WARNING: No 'on_success' set for {:?}, default of 'OnSuccess::Continue'", name);
                        OnSuccess::Continue
                    },
                    Some(toml_value) => {
                        let table = toml_value.as_table().unwrap();
                        let (message_type, message) = table.iter().nth(0).unwrap();
                        match message_type.as_ref() {
                            "echo" => OnSuccess::Echo {message: message.as_str().unwrap().to_string()},
                            "warn" => OnSuccess::Warn {message: message.as_str().unwrap().to_string()},
                            "abort" => OnSuccess::Abort {message: message.as_str().unwrap().to_string()},
                            _ => unimplemented!()
                        }
                    }
                };

                let on_failure = match details.get("on_failure") {
                    None => {
                        println!("WARNING: No 'on_failure' set for {:?}, default of 'OnFailure::Continue'", name);
                        OnFailure::Continue
                    },
                    Some(toml_value) => {
                        let table = toml_value.as_table().unwrap();
                        let (message_type, message) = table.iter().nth(0).unwrap();
                        match message_type.as_ref() {
                            "echo" => OnFailure::Echo {message: message.as_str().unwrap().to_string()},
                            "warn" => OnFailure::Warn {message: message.as_str().unwrap().to_string()},
                            "abort" => OnFailure::Abort {message: message.as_str().unwrap().to_string()},
                            _ => unimplemented!()
                        }
                    }
                };

                Step::Run {name: name, command: command, arguments: arguments, on_success: on_success, on_failure: on_failure}
            },
            "setenv" => {
                let variable = details.get("variable").unwrap().as_str().unwrap().to_string();
                let value = details.get("value").unwrap().as_str().unwrap().to_string();
                Step::SetEnv {variable: variable, value: value}
            },
            "shell" => {
                let command = details.get("command").unwrap().as_str().unwrap().to_string();

                let stdout = match details.get("stdout") {
                    None => {
                        Pipe::StdOut
                    },
                    Some(toml_value) => {
                        Pipe::Variable {name: toml_value.as_str().unwrap().to_owned()}
                    }
                };

                let on_success = match details.get("on_success") {
                    None => {
                        println!("WARNING: No 'on_success' set for {:?}, default of 'OnSuccess::Continue'", name);
                        OnSuccess::Continue
                    },
                    Some(toml_value) => {
                        let table = toml_value.as_table().unwrap();
                        let (message_type, message) = table.iter().nth(0).unwrap();
                        match message_type.as_ref() {
                            "echo" => OnSuccess::Echo {message: message.as_str().unwrap().to_string()},
                            "warn" => OnSuccess::Warn {message: message.as_str().unwrap().to_string()},
                            "abort" => OnSuccess::Abort {message: message.as_str().unwrap().to_string()},
                            _ => unimplemented!()
                        }
                    }
                };

                let on_failure = match details.get("on_failure") {
                    None => {
                        println!("WARNING: No 'on_failure' set for {:?}, default of 'OnFailure::Continue'", name);
                        OnFailure::Continue
                    },
                    Some(toml_value) => {
                        let table = toml_value.as_table().unwrap();
                        let (message_type, message) = table.iter().nth(0).unwrap();
                        match message_type.as_ref() {
                            "echo" => OnFailure::Echo {message: message.as_str().unwrap().to_string()},
                            "warn" => OnFailure::Warn {message: message.as_str().unwrap().to_string()},
                            "abort" => OnFailure::Abort {message: message.as_str().unwrap().to_string()},
                            _ => unimplemented!()
                        }
                    }
                };

                Step::Shell {name: name, command: command, stdout: stdout, on_success: on_success, on_failure: on_failure}
            },
            _ => {
                println!("Unknown action '{}' in toml file", action);
                unimplemented!();
            },
        };
        // println!("Step: {:?}", step);

        steps.push(step);
    }

    steps
}

pub fn execute_steps(filename: &str) {

    let steps = parse_toml(filename);

    println!("Steps: {:#?}", steps);

    // Loop over steps
    let nb_steps = steps.len();
    for (idx, step) in steps.into_iter().enumerate() {
        // println!("Step {}/{} -- ", idx+1, nb_steps);
        match step {
            Step::Echo {name, message} => {
                // NOTE: If an environment variable is non-existing, the following line will panic.
                println!("{}", shellexpand::full(&message).unwrap());
            },
            Step::Run {name, command, arguments, on_success, on_failure} => {
                // NOTE: If an environment variable is non-existing, the following line will panic.
                let command: String = shellexpand::full(&command).unwrap().to_string();
                // NOTE: If an environment variable is non-existing, the following line will panic.
                let arguments: String = shellexpand::full(&arguments).unwrap().to_string();
                println!("    command: {} {}", command, arguments);
                let arguments: Vec<&str> = arguments.split(' ').collect();
                // println!("    command: {}   arguments: {:?}", command, arguments);
                let output = std::process::Command::new(command)
                                                       .args(&arguments)
                                                       .output()
                                                       .unwrap_or_else(|e| { panic!("failed to execute process: {}", e) });

                // println!("output: {:#?}", output);

                // Print stdout
                println!("{}", String::from_utf8_lossy(&output.stdout));
                println!("{}", String::from_utf8_lossy(&output.stderr));

                if output.status.success() {
                    match on_success {
                        OnSuccess::Continue => {},
                        OnSuccess::Echo {message} => {
                            println!("{}", message);
                        },
                        OnSuccess::Warn {message} => {
                            println!("WARNING: {}", message);
                        },
                        OnSuccess::Abort {message} => {
                            println!("ABORT: {}", message);
                            break;
                        },
                    }
                } else {
                    match on_failure {
                        OnFailure::Continue => {},
                        OnFailure::Echo {message} => {
                            println!("{}", message);
                        },
                        OnFailure::Warn {message} => {
                            println!("WARNING: {}", message);
                        },
                        OnFailure::Abort {message} => {
                            println!("ABORT: {}", message);
                            break;
                        },
                    }
                }
            },
            Step::SetEnv {variable, value} => {
                std::env::set_var(variable, value);
            },
            Step::Shell {name, command, stdout, on_success, on_failure} => {
                let shell = "sh";
                println!("    command: {} -c '{}'", shell, command);
                let output = std::process::Command::new(shell)
                                                       .arg("-c")
                                                       .arg(command)
                                                       .output()
                                                       .unwrap_or_else(|e| { panic!("failed to execute process: {}", e) });

               // Print stdout
                match stdout {
                    Pipe::Null => { },
                    Pipe::StdOut => {
                        println!("{}", String::from_utf8_lossy(&output.stdout));
                    },
                    Pipe::StdErr => {
                        // FIXME: Outout to stderr, not stdout
                        unimplemented!();
                    },
                    Pipe::Variable {name} => {
                        std::env::set_var(name, String::from_utf8_lossy(&output.stdout).into_owned());
                    },
                    Pipe::File {filename} => {
                        // FIXME: Save to file
                        unimplemented!();
                    }
                };

               if output.status.success() {
                   match on_success {
                       OnSuccess::Continue => {},
                       OnSuccess::Echo {message} => {
                           println!("{}", message);
                       },
                       OnSuccess::Warn {message} => {
                           println!("WARNING: {}", message);
                       },
                       OnSuccess::Abort {message} => {
                           println!("ABORT: {}", message);
                           break;
                       },
                   }
               } else {
                   match on_failure {
                       OnFailure::Continue => {},
                       OnFailure::Echo {message} => {
                           println!("{}", message);
                       },
                       OnFailure::Warn {message} => {
                           println!("WARNING: {}", message);
                       },
                       OnFailure::Abort {message} => {
                           println!("ABORT: {}", message);
                           break;
                       },
                   }
               }
            },
        }
    }

}
<commit_msg>Expand some variables before printing<commit_after>extern crate toml;
extern crate shellexpand;

// Import traits
use std::io::prelude::Read;

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}

pub fn test() {
    println!("Hello from lib.rs");
}

#[derive(Debug)]
enum OnSuccess {
    Continue,
    Echo {message: String},
    Warn {message: String},
    Abort {message: String},
}

#[derive(Debug)]
enum OnFailure {
    Continue,
    Echo {message: String},
    Warn {message: String},
    Abort {message: String},
}

#[derive(Debug)]
enum Pipe {
    Null,
    StdOut,
    StdErr,
    Variable {name: String},
    File {filename: String},
}

#[derive(Debug)]
enum Step {
    Echo {name: String, message: String},
    Run {name: String, command: String, arguments: String, on_success: OnSuccess, on_failure: OnFailure},
    Shell {name: String, command: String, stdout: Pipe, on_success: OnSuccess, on_failure: OnFailure},
    SetEnv {variable: String, value: String},
}

fn parse_toml(filename: &str) -> Vec<Step> {

    println!("Reading toml file {}", filename);

    // Read toml file content
    let mut input = String::new();
    std::fs::File::open(&filename).and_then(|mut f| {
        f.read_to_string(&mut input)
    }).unwrap();

    // Pase toml file
    let mut parser = toml::Parser::new(&input);

    let toml = parser.parse().unwrap_or_else(|| {
        for err in &parser.errors {
            let (loline, locol) = parser.to_linecol(err.lo);
            let (hiline, hicol) = parser.to_linecol(err.hi);
            println!("{}:{}:{}-{}:{} error: {}",
                     filename, loline, locol, hiline, hicol, err.desc);
        }
        panic!("Error parsing toml file.");
    });

    // println!("toml: {:#?}", toml);

    // Parse "variables" table
    for (variable, value) in toml.get("variables").unwrap().as_table().unwrap() {
        let value = value.as_str().unwrap();
        // println!("variable: {}   value: {}", variable, value);
        std::env::set_var(variable, value);
    }

    // Parse "steps" table
    let mut steps = Vec::<Step>::with_capacity(10);

    for step in toml.get("step").unwrap().as_slice().unwrap() {
        let action = step.as_table().unwrap().get("action").unwrap().as_str().unwrap();
        let name = step.as_table().unwrap().get("name").unwrap().as_str().unwrap().to_string();
        let details = step.as_table().unwrap().get("details").unwrap().as_table().unwrap();
        let step = match action {
            "echo" => {
                let message = details.get("message").unwrap().as_str().unwrap().to_string();
                Step::Echo {name: name, message: message}
            },
            "run" => {
                let command = details.get("command").unwrap().as_str().unwrap().to_string();

                // Allow non-existing 'arguments' field, defaulting to empty string.
                let arguments_option = details.get("arguments");
                let arguments = match arguments_option {
                    None => String::new(),
                    Some(arguments_toml_value) => arguments_toml_value.as_str().unwrap().to_string(),
                };

                let on_success = match details.get("on_success") {
                    None => {
                        println!("WARNING: No 'on_success' set for {:?}, default of 'OnSuccess::Continue'", name);
                        OnSuccess::Continue
                    },
                    Some(toml_value) => {
                        let table = toml_value.as_table().unwrap();
                        let (message_type, message) = table.iter().nth(0).unwrap();
                        match message_type.as_ref() {
                            "echo" => OnSuccess::Echo {message: message.as_str().unwrap().to_string()},
                            "warn" => OnSuccess::Warn {message: message.as_str().unwrap().to_string()},
                            "abort" => OnSuccess::Abort {message: message.as_str().unwrap().to_string()},
                            _ => unimplemented!()
                        }
                    }
                };

                let on_failure = match details.get("on_failure") {
                    None => {
                        println!("WARNING: No 'on_failure' set for {:?}, default of 'OnFailure::Continue'", name);
                        OnFailure::Continue
                    },
                    Some(toml_value) => {
                        let table = toml_value.as_table().unwrap();
                        let (message_type, message) = table.iter().nth(0).unwrap();
                        match message_type.as_ref() {
                            "echo" => OnFailure::Echo {message: message.as_str().unwrap().to_string()},
                            "warn" => OnFailure::Warn {message: message.as_str().unwrap().to_string()},
                            "abort" => OnFailure::Abort {message: message.as_str().unwrap().to_string()},
                            _ => unimplemented!()
                        }
                    }
                };

                Step::Run {name: name, command: command, arguments: arguments, on_success: on_success, on_failure: on_failure}
            },
            "setenv" => {
                let variable = details.get("variable").unwrap().as_str().unwrap().to_string();
                let value = details.get("value").unwrap().as_str().unwrap().to_string();
                Step::SetEnv {variable: variable, value: value}
            },
            "shell" => {
                let command = details.get("command").unwrap().as_str().unwrap().to_string();

                let stdout = match details.get("stdout") {
                    None => {
                        Pipe::StdOut
                    },
                    Some(toml_value) => {
                        Pipe::Variable {name: toml_value.as_str().unwrap().to_owned()}
                    }
                };

                let on_success = match details.get("on_success") {
                    None => {
                        println!("WARNING: No 'on_success' set for {:?}, default of 'OnSuccess::Continue'", name);
                        OnSuccess::Continue
                    },
                    Some(toml_value) => {
                        let table = toml_value.as_table().unwrap();
                        let (message_type, message) = table.iter().nth(0).unwrap();
                        match message_type.as_ref() {
                            "echo" => OnSuccess::Echo {message: message.as_str().unwrap().to_string()},
                            "warn" => OnSuccess::Warn {message: message.as_str().unwrap().to_string()},
                            "abort" => OnSuccess::Abort {message: message.as_str().unwrap().to_string()},
                            _ => unimplemented!()
                        }
                    }
                };

                let on_failure = match details.get("on_failure") {
                    None => {
                        println!("WARNING: No 'on_failure' set for {:?}, default of 'OnFailure::Continue'", name);
                        OnFailure::Continue
                    },
                    Some(toml_value) => {
                        let table = toml_value.as_table().unwrap();
                        let (message_type, message) = table.iter().nth(0).unwrap();
                        match message_type.as_ref() {
                            "echo" => OnFailure::Echo {message: message.as_str().unwrap().to_string()},
                            "warn" => OnFailure::Warn {message: message.as_str().unwrap().to_string()},
                            "abort" => OnFailure::Abort {message: message.as_str().unwrap().to_string()},
                            _ => unimplemented!()
                        }
                    }
                };

                Step::Shell {name: name, command: command, stdout: stdout, on_success: on_success, on_failure: on_failure}
            },
            _ => {
                println!("Unknown action '{}' in toml file", action);
                unimplemented!();
            },
        };
        // println!("Step: {:?}", step);

        steps.push(step);
    }

    steps
}

pub fn execute_steps(filename: &str) {

    let steps = parse_toml(filename);

    println!("Steps: {:#?}", steps);

    // Loop over steps
    let nb_steps = steps.len();
    for (idx, step) in steps.into_iter().enumerate() {
        // println!("Step {}/{} -- ", idx+1, nb_steps);
        match step {
            Step::Echo {name, message} => {
                // NOTE: If an environment variable is non-existing, the following line will panic.
                println!("{}", shellexpand::full(&message).unwrap());
            },
            Step::Run {name, command, arguments, on_success, on_failure} => {
                // NOTE: If an environment variable is non-existing, the following line will panic.
                let command: String = shellexpand::full(&command).unwrap().to_string();
                // NOTE: If an environment variable is non-existing, the following line will panic.
                let arguments: String = shellexpand::full(&arguments).unwrap().to_string();
                println!("    command: {} {}", command, arguments);
                let arguments: Vec<&str> = arguments.split(' ').collect();
                // println!("    command: {}   arguments: {:?}", command, arguments);
                let output = std::process::Command::new(command)
                                                       .args(&arguments)
                                                       .output()
                                                       .unwrap_or_else(|e| { panic!("failed to execute process: {}", e) });

                // println!("output: {:#?}", output);

                // Print stdout
                println!("{}", String::from_utf8_lossy(&output.stdout));
                println!("{}", String::from_utf8_lossy(&output.stderr));

                if output.status.success() {
                    match on_success {
                        OnSuccess::Continue => {},
                        OnSuccess::Echo {message} => {
                            println!("{}", shellexpand::full(&message).unwrap());
                        },
                        OnSuccess::Warn {message} => {
                            println!("WARNING: {}", shellexpand::full(&message).unwrap());
                        },
                        OnSuccess::Abort {message} => {
                            println!("ABORT: {}", shellexpand::full(&message).unwrap());
                            break;
                        },
                    }
                } else {
                    match on_failure {
                        OnFailure::Continue => {},
                        OnFailure::Echo {message} => {
                            println!("{}", shellexpand::full(&message).unwrap());
                        },
                        OnFailure::Warn {message} => {
                            println!("WARNING: {}", shellexpand::full(&message).unwrap());
                        },
                        OnFailure::Abort {message} => {
                            println!("ABORT: {}", shellexpand::full(&message).unwrap());
                            break;
                        },
                    }
                }
            },
            Step::SetEnv {variable, value} => {
                std::env::set_var(variable, value);
            },
            Step::Shell {name, command, stdout, on_success, on_failure} => {
                let shell = "sh";
                println!("    command: {} -c '{}'", shell, command);
                let output = std::process::Command::new(shell)
                                                       .arg("-c")
                                                       .arg(command)
                                                       .output()
                                                       .unwrap_or_else(|e| { panic!("failed to execute process: {}", e) });

               // Print stdout
                match stdout {
                    Pipe::Null => { },
                    Pipe::StdOut => {
                        println!("{}", String::from_utf8_lossy(&output.stdout));
                    },
                    Pipe::StdErr => {
                        // FIXME: Outout to stderr, not stdout
                        unimplemented!();
                    },
                    Pipe::Variable {name} => {
                        std::env::set_var(name, String::from_utf8_lossy(&output.stdout).into_owned());
                    },
                    Pipe::File {filename} => {
                        // FIXME: Save to file
                        unimplemented!();
                    }
                };

               if output.status.success() {
                   match on_success {
                       OnSuccess::Continue => {},
                       OnSuccess::Echo {message} => {
                           println!("{}", shellexpand::full(&message).unwrap());
                       },
                       OnSuccess::Warn {message} => {
                           println!("WARNING: {}", shellexpand::full(&message).unwrap());
                       },
                       OnSuccess::Abort {message} => {
                           println!("ABORT: {}", shellexpand::full(&message).unwrap());
                           break;
                       },
                   }
               } else {
                   match on_failure {
                       OnFailure::Continue => {},
                       OnFailure::Echo {message} => {
                           println!("{}", shellexpand::full(&message).unwrap());
                       },
                       OnFailure::Warn {message} => {
                           println!("WARNING: {}", shellexpand::full(&message).unwrap());
                       },
                       OnFailure::Abort {message} => {
                           println!("ABORT: {}", shellexpand::full(&message).unwrap());
                           break;
                       },
                   }
               }
            },
        }
    }

}
<|endoftext|>