<commit_before>//! DFQ is a decoupled, fault-tolerant, multi-producer single-consumer queue.
//! DFQ is compatible with `no_std` and is optionally interrupt-safe by being "temporarily lockless".
//!
//! DFQ accepts immutable data items only and does not allow the consumer
//! to remove (pop) or modify the items in the queue.
//! Instead, the consumer can call `mark_completed()` on a queued item
//! to indicate to the producer that the item is ready to be removed from the queue.
//!
//! Each producer retains ownership of the data items it queues, and only those items.
//! Thus, a producer is responsible for removing its items from the queue once they have been marked complete by the consumer.

// #![no_std]

#![allow(dead_code)]
#![feature(alloc, collections)]

extern crate core;

extern crate spin;
extern crate alloc;
extern crate collections;

use spin::{Mutex, MutexGuard};
use core::sync::atomic::{AtomicBool, Ordering};
use alloc::arc::Arc;
use collections::VecDeque;

/// These two types indicate that only immutable items can be queued right now.
pub type EnqueueParam<T> = Arc<T>;
type QueuedDataType<T> = Arc<QueuedData<T>>;

/// A wrapper around the data that has been queued. This data type is returned when peeking into the queue.
/// private note:  This must NOT be cloned because it contains an inner Arc reference to the data passed to its constructor!
#[derive(Debug)]
pub struct QueuedData<T> {
    data: Arc<T>,
    completed: AtomicBool,
}

impl<T> QueuedData<T> {
    /// this does not take ownership of the given `dataref` Arc instance
    /// (nor does it borrow it long term),
    /// allowing you as the caller to retain ownership of that data.
    fn new(dataref: Arc<T>) -> QueuedData<T> {
        QueuedData {
            data: dataref.clone(),
            completed: AtomicBool::new(false),
        }
    }

    /// Mark this item as completed.
    /// This allows the producer (the one that pushed it onto the queue) to know that it's safe to remove it from the queue.
    pub fn mark_completed(&self) {
        self.completed.store(true, Ordering::SeqCst);
    }

    /// Whether this item has been completed (handled) by the DFQueueConsumer.
    /// If an item is completed, the producer knows it's okay to remove it from the queue.
    pub fn is_completed(&self) -> bool {
        self.completed.load(Ordering::SeqCst)
    }

    // /// returns true if the given `QueuedData` was locked
    // fn is_locked(&self) -> bool {
    //     match self.data.try_lock() {
    //         Some(_) => false,
    //         None => true,
    //     }
    // }
}

/// taken from Rust's VecDeque implementation
const INITIAL_CAPACITY: usize = 7; // 2^3 - 1

/// The actual queue, an opaque type that cannot be used directly.
/// The user must use `DFQueueConsumer` and `DFQueueProducer`.
#[derive(Debug)]
pub struct DFQueue<T> {
    /// the actual inner queue
    queue: Mutex<VecDeque<QueuedDataType<T>>>,
    /// whether this queue has a consumer (it can only have one!!)
    has_consumer: AtomicBool,
}

impl<T> DFQueue<T> {

    /// Creates a new DFQueue.
    ///
    /// This object cannot be used directly, you must obtain a producer or consumer to the queue
    /// using the functions `into_consumer()` or `obtain_producer()`.
    pub fn new() -> DFQueue<T> {
        DFQueue::with_capacity(INITIAL_CAPACITY)
    }

    /// Creates a new DFQueue with the given initial_capacity.
    ///
    /// This object cannot be used directly, you must obtain a producer or consumer to the queue
    /// using the functions `into_consumer()` or `obtain_producer()`.
    pub fn with_capacity(initial_capacity: usize) -> DFQueue<T> {
        DFQueue {
            queue: Mutex::new(VecDeque::with_capacity(initial_capacity)),
            has_consumer: AtomicBool::new(false),
        }
    }

    /// Consumes the DFQueue and returns the one and only consumer for this DFQueue.
    /// It consumes the DFQueue instance because there is only one consumer allowed per DFQueue.
    pub fn into_consumer(self) -> DFQueueConsumer<T> {
        debug_assert!(self.has_consumer.load(Ordering::SeqCst), "DFQueue::into_consumer(): WTF? already had a consumer!");
        self.has_consumer.store(true, Ordering::SeqCst);

        DFQueueConsumer {
            qref: Arc::new(self),
        }
    }

    /// Consumes the DFQueue and returns a producer.
    /// To obtain another DFQueueProducer for this DFQueue, call `obtain_producer()` on the returned DFQueueProducer.
    /// DFQueueProducer does not implement the standard Clone trait, to avoid accidentally cloning it implicity.
    pub fn into_producer(self) -> DFQueueProducer<T> {
        DFQueueProducer {
            qref: Arc::new(self),
            local_backlog: VecDeque::new(), // TODO: may need to wrap this in an Rc<>
        }
    }

}

/// A consumer that can process (peek into) elements in a DFQueue, but not actually remove them.
/// Do not wrap this in an Arc or Mutex, the queue it is already protected by those on the interior.
///
/// This does not provide a `pop()` method like most queues,
/// because we do not permit the consumer to remove items from the queue.
/// Instead, we require that an element can only be removed from the queue by the producer who originally enqueued it.
#[derive(Debug)]
pub struct DFQueueConsumer<T> {
    qref: Arc<DFQueue<T>>,
}

impl<T> DFQueueConsumer<T> {

    /// obtains a DFQueueProducer from this consumer instance, since there can be multiple producers.
    pub fn obtain_producer(&self) -> DFQueueProducer<T> {
        DFQueueProducer {
            qref: self.qref.clone(),
            local_backlog: VecDeque::new(), // TODO: may need to wrap this in an Rc<>
        }
    }

    /// internal function that returns the first non-completed item in the queue (queue is locked already)
    fn peek_locked(&self, locked_queue: MutexGuard<VecDeque<QueuedDataType<T>>>) -> Option<QueuedDataType<T>> {
        for e in locked_queue.iter() {
            if e.completed.load(Ordering::SeqCst) {
                continue;
            }
            else {
                return Some(e.clone());
            }
        }

        None
    }

    /// Peeks at the queue in a blocking fashion.
    /// This will block until it can acquire a lock on the inner queue.
    /// Returns the first non-completed element in the queue without actually removing it from the queue,
    /// or `None` if the queue is empty.
    ///
    /// Note: this is NOT IRQ-safe.
    /// TODO: use a data guard to automatically release the lock to the returned QueuedData
    pub fn peek_locking(&self) -> Option<QueuedDataType<T>> {
        // first, acquire the lock so we have exclusive access to the inner queue
        let innerq = self.qref.queue.lock(); // blocking call

        self.peek_locked(innerq)
    }

    /// Peeks at the queue in a non-blocking fashion.
    /// Returns the last non-completed element in the queue without actually removing it from the queue.
    /// Returns `None` if the lock to the inner queue cannot be obtained, or if the queue is empty.
    ///
    /// Note: this is IRQ-safe.
    /// TODO: use a data guard to automatically release the lock to the returned QueuedData
    pub fn peek_lockless(&self) -> Option<QueuedDataType<T>> {
        // first, try to acquire the queue lock.
        let lock_result = self.qref.queue.try_lock();
        match lock_result {
            Some(innerq) => {
                self.peek_locked(innerq)
            },
            None => {
                None
            }
        }
    }

}

/// A producer that can enqueue elements into a DFQueue.
/// Do not wrap this in an Arc or Mutex, the queue it is already protected by those on the interior.
#[derive(Debug)]
pub struct DFQueueProducer<T> {
    qref: Arc<DFQueue<T>>,
    local_backlog: VecDeque<QueuedDataType<T>>,
}

impl<T> DFQueueProducer<T> {

    /// Call this to obtain another DFQueueProducer.
    /// DFQueueProducer does not implement the standard Clone trait, to avoid accidentally cloning it implicity.
    pub fn obtain_producer(&self) -> DFQueueProducer<T> {
        DFQueueProducer {
            qref: self.qref.clone(),
            local_backlog: VecDeque::new(), // TODO: may need to wrap this in an Rc<>
        }
    }

    /// Returns a DFQueueConsumer for this queue, if it hasn't yet been obtained
    /// (either via this function or via `DFQueue::into_consumer()`).
    /// To ensure there is only a single DFQueueConsumer, it will return `None` if there is already a `DFQueueConsumer`.
    pub fn get_consumer(&self) -> Option<DFQueueConsumer<T>> {
        let has_consumer: bool = self.qref.has_consumer.load(Ordering::SeqCst);
        match has_consumer {
            true => None,
            false => {
                self.qref.has_consumer.store(true, Ordering::SeqCst);
                Some(
                    DFQueueConsumer {
                        qref: self.qref.clone(),
                    }
                )
            }
        }
    }

    /// pushes the given `QueuedData`, which is an Arc reference wrapping the actual data item, onto the queue.
    /// This is a blocking call that will acquire the queue lock, write the data,
    /// and release the lock before returning to the caller.
    ///
    /// Note: this is NOT IRQ-safe.
    pub fn enqueue_locking(&mut self, dataref: EnqueueParam<T>) {

        // acquire the lock so we have exclusive access to the inner queue
        let mut innerq = self.qref.queue.lock(); // blocking call

        // 1) move all the previously-backlogged data into the inner queue
        // 2) push the actual `dataref` onto the queue
        innerq.append(&mut self.local_backlog);
        innerq.push_back(Arc::new(QueuedData::new(dataref.clone())));

    }

    /// pushes the given `QueuedData`, which is an Arc reference wrapping the actual data item,
    /// onto the back of the queue, only if the queue lock can be acquired immediately.
    /// If the queue lock cannot be acquired, then the given `QueuedData` is put on a local backlog queue
    /// that is specific to this `DFQueueRef`, which will be merged into the actual inner queue
    /// on the next invocation of any `enqueue` method, provided that the inner queue can be locked immediately at that point.
    /// Returns true if the given `dataref` was added to the inner queue directly, false if it was added to the backlog.
    ///
    /// Note: this is IRQ-safe.
    pub fn enqueue_lockless(&mut self, dataref: EnqueueParam<T>) -> bool {

        // first, try to acquire the queue lock.
        match self.qref.queue.try_lock() {
            Some(mut innerq) => {

                // 1) move all the previously-backlogged data into the inner queue
                // 2) push the actual `dataref` onto the queue
                innerq.append(&mut self.local_backlog);
                innerq.push_back(Arc::new(QueuedData::new(dataref.clone())));
                true
            },
            None => {
                // if we can NOT acquire it:
                // add the given `dataref` to this DFQueueRef's producer-local backlog queue
                self.local_backlog.push_back(Arc::new(QueuedData::new(dataref.clone())));
                false
            }
        }

    }

}

// Conditionally compile the module `test` only when the test-suite is run.
#[cfg(test)]
mod test {

    extern crate std;

    use std::thread;
    use super::*;

    #[test]
    // #[should_panic]
    fn simple_test() {

        let queue: DFQueue<usize> = DFQueue::new();

        let mut queue_prod = queue.into_producer();
        let mut queue_cons = queue_prod.get_consumer();
        assert!(queue_cons.is_some(), "First DFQueueConsumer was None!!!");

        // a second call to get_consumer must return None
        assert!(queue_prod.get_consumer().is_none(), "Second DFQueueConsumer wasn't None!!");

        let queue_cons = queue_cons.unwrap();

        let mut queue_prod2 = queue_prod.obtain_producer();

        let mut thr_p = thread::spawn( move || {
            let original_data: Vec<Arc<usize>> = vec![Arc::new(1), Arc::new(2), Arc::new(3), Arc::new(4), Arc::new(5)];

            for i in 1..20 {
                for elem in original_data.iter() {
                    queue_prod.enqueue_lockless(elem.clone());
                }

                let queued_data: Arc<usize> = Arc::new(255);
                queue_prod.enqueue_lockless(queued_data);
            }
        } );

         let mut thr_p2 = thread::spawn( move || {
            let original_data: Vec<Arc<usize>> = vec![Arc::new(10), Arc::new(11), Arc::new(12), Arc::new(13), Arc::new(14)];

            for i in 1..20 {
                for elem in original_data.iter() {
                    queue_prod2.enqueue_lockless(elem.clone());
                }

                let queued_data: Arc<usize> = Arc::new(512);
                queue_prod2.enqueue_lockless(queued_data);
            }
        } );

        let mut thr_c = thread::spawn( move || {
            loop {
                let mut val = queue_cons.peek_locking();
                if let Some(v) = val {
                    println!("peeked: {:?}", v);
                    v.mark_completed();
                }
            }
        } );

        thr_p.join();
        thr_p2.join();
        thr_c.join();
    }
}<commit_msg>Separated the returned values for QueuedData (for the producers) and PeekedData (for the consumer).<commit_after>//! DFQ is a decoupled, fault-tolerant, multi-producer single-consumer queue.
//! DFQ is compatible with `no_std` and is optionally interrupt-safe by being "temporarily lockless".
//!
//! DFQ accepts immutable data items only and does not allow the consumer
//! to remove (pop) or modify the items in the queue.
//! Instead, the consumer can call `mark_completed()` on a queued item
//! to indicate to the producer that the item is ready to be removed from the queue.
//!
//! Each producer retains ownership of the data items it queues, and only those items.
//! Thus, a producer is responsible for removing its items from the queue once they have been marked complete by the consumer.

// #![no_std]

#![allow(dead_code)]
#![feature(alloc, collections)]

extern crate core;

extern crate spin;
extern crate alloc;
extern crate collections;

use spin::{Mutex, MutexGuard};
use core::sync::atomic::{AtomicBool, Ordering};
use core::ops::Deref;
use alloc::arc::Arc;
use collections::VecDeque;

/// A special reference type that wraps a data item that has been queued.
/// This is returned to a producer thread (the user of a DFQueueProducer)
/// when enqueuing an item onto the queue so that the producer
/// can retain a reference to it in the case of failure.
#[derive(Debug)]
pub struct QueuedData<T> {
    inner: Arc<InnerQueuedData<T>>,
}

impl<T> QueuedData<T> {

    /// Not public, a DFQueueProducer must call an enqueue function to receive one of these back.
    fn new(data: T) -> QueuedData<T> {
        QueuedData {
            inner: Arc::new(InnerQueuedData::new(data)),
        }
    }

    /// Whether this item has been completed (handled) by the DFQueueConsumer.
    /// If an item is completed, the producer knows it's okay to remove it from the queue.
    pub fn is_completed(&self) -> bool {
        self.inner.completed.load(Ordering::SeqCst)
    }

    /// Returns true if this data is still on the queue
    ///
    /// The logic here is as follows: the thread invoking this function holds one reference.
    /// Thus, if there is more than one reference, then it means it is on the queue,
    /// because the queue also holds a reference to it.
    /// That's why we cannot call it internally --  it won't be correct because the original producer thread
    /// may also be holding a reference to it, in order for that producer to retain a reference to it
    ////
    ///
    /// Private note: do not call this internally! This is a public API meant for the producer thread to use.
    pub fn is_enqueued(&self) -> bool {
        Arc::strong_count(&self.inner) > 1
    }

    /// The logic here is as follows: if the item on the queue has only one reference, and has not been completed,
    /// that means it is no longer on the queue (the consumer thread crashed or failed).
    /// Again, the producer thread invoking this function holds one reference, and if the data is indeed enqueued,
    /// that constitutes another reference.
    /// If the reference held in the queue isn't there (only one total), and the thread hasn't completed
    /// (it may have just been completed in another enqueue call), then something went wrong because it's not complete
    /// and not on the queue anymore.
    ///
    /// Private note: do not call this internally! This is a public API meant for the producer thread to use.
    pub fn has_failed(&self) -> bool {
        // an item that has been completed could not have possibly failed!
        (!self.is_completed()) &&
        (Arc::strong_count(&self.inner) == 1)
    }

    /// creates a clone of this, marshalled as a PeekedData
    /// so that a consumer can access the InnerQueuedData without being able to access
    /// all of the QueuedData methods here.
    fn as_peeked(&self) -> PeekedData<T> {
        PeekedData {
            inner: self.inner.clone(),
        }
    }

    /// not public, because we want to control when this is cloned
    /// so we can keep track of the number of Arc references to the InnerQueuedData.
    fn clone(&self) -> QueuedData<T> {
        QueuedData {
            inner: self.inner.clone(),
        }
    }
}

#[derive(Debug)]
struct InnerQueuedData<T> {
    data: T,
    completed: AtomicBool,
}

impl<T> InnerQueuedData<T> {

    fn new(data: T) -> InnerQueuedData<T> {
        InnerQueuedData {
            data: data,
            completed: AtomicBool::new(false),
        }
    }

}

/// A wrapper around data in the queue that allows a DFQueueConsumer
/// to access the data and mark the queued item as completed.
/// Automatically Derefs to the inner type `&T`, just like Arc does.
#[derive(Debug)]
pub struct PeekedData<T> {
    inner: Arc<InnerQueuedData<T>>,
}

impl<T> PeekedData<T> {

    /// Mark this item as completed.
    /// This allows the producer (the one that pushed it onto the queue) to know that it's safe to remove it from the queue.
    pub fn mark_completed(&self) {
        self.inner.completed.store(true, Ordering::SeqCst);
    }

}

impl<T> Deref for PeekedData<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.inner.data
    }
}

/// taken from Rust's VecDeque implementation
const INITIAL_CAPACITY: usize = 7; // 2^3 - 1

/// The actual queue, an opaque type that cannot be used directly.
/// The user must use `DFQueueConsumer` and `DFQueueProducer`.
#[derive(Debug)]
pub struct DFQueue<T> {
    /// the actual inner queue
    queue: Mutex<VecDeque<QueuedData<T>>>,
    /// whether this queue has a consumer (it can only have one!!)
    has_consumer: AtomicBool,
}

impl<T> DFQueue<T> {

    /// Creates a new DFQueue.
    ///
    /// This object cannot be used directly, you must obtain a producer or consumer to the queue
    /// using the functions `into_consumer()` or `obtain_producer()`.
    pub fn new() -> DFQueue<T> {
        DFQueue::with_capacity(INITIAL_CAPACITY)
    }

    /// Creates a new DFQueue with the given initial_capacity.
    ///
    /// This object cannot be used directly, you must obtain a producer or consumer to the queue
    /// using the functions `into_consumer()` or `obtain_producer()`.
    pub fn with_capacity(initial_capacity: usize) -> DFQueue<T> {
        DFQueue {
            queue: Mutex::new(VecDeque::with_capacity(initial_capacity)),
            has_consumer: AtomicBool::new(false),
        }
    }

    /// Consumes the DFQueue and returns the one and only consumer for this DFQueue.
    /// It consumes the DFQueue instance because there is only one consumer allowed per DFQueue.
    pub fn into_consumer(self) -> DFQueueConsumer<T> {
        debug_assert!(self.has_consumer.load(Ordering::SeqCst), "DFQueue::into_consumer(): WTF? already had a consumer!");
        self.has_consumer.store(true, Ordering::SeqCst);

        DFQueueConsumer {
            qref: Arc::new(self),
        }
    }

    /// Consumes the DFQueue and returns a producer.
    /// To obtain another DFQueueProducer for this DFQueue, call `obtain_producer()` on the returned DFQueueProducer.
    /// DFQueueProducer does not implement the standard Clone trait, to avoid accidentally cloning it implicity.
    pub fn into_producer(self) -> DFQueueProducer<T> {
        DFQueueProducer {
            qref: Arc::new(self),
            local_backlog: VecDeque::new(), // TODO: may need to wrap this in an Rc<>
        }
    }

}

/// A consumer that can process (peek into) elements in a DFQueue, but not actually remove them.
/// Do not wrap this in an Arc or Mutex, the queue it is already protected by those on the interior.
///
/// This does not provide a `pop()` method like most queues,
/// because we do not permit the consumer to remove items from the queue.
/// Instead, we require that an element can only be removed from the queue by the producer who originally enqueued it.
#[derive(Debug)]
pub struct DFQueueConsumer<T> {
    qref: Arc<DFQueue<T>>,
}

impl<T> DFQueueConsumer<T> {

    /// obtains a DFQueueProducer from this consumer instance, since there can be multiple producers.
    pub fn obtain_producer(&self) -> DFQueueProducer<T> {
        DFQueueProducer {
            qref: self.qref.clone(),
            local_backlog: VecDeque::new(),
        }
    }

    /// internal function that returns the first non-completed item in the queue (queue is locked already)
    fn peek_locked(&self, locked_queue: MutexGuard<VecDeque<QueuedData<T>>>) -> Option<PeekedData<T>> {
        for e in locked_queue.iter() {
            if e.is_completed() {
                continue;
            }
            else {
                return Some(e.as_peeked());
            }
        }

        None
    }

    /// Peeks at the queue in a blocking fashion.
    /// This will block until it can acquire a lock on the inner queue.
    /// Returns the first non-completed element in the queue without actually removing it from the queue,
    /// or `None` if the queue is empty.
    ///
    /// Note: this is NOT IRQ-safe.
    /// TODO: use a data guard to automatically release the lock to the returned InnerQueuedData
    pub fn peek_locking(&self) -> Option<PeekedData<T>> {
        // first, acquire the lock so we have exclusive access to the inner queue
        let innerq = self.qref.queue.lock(); // blocking call

        self.peek_locked(innerq)
    }

    /// Peeks at the queue in a non-blocking fashion.
    /// Returns the last non-completed element in the queue without actually removing it from the queue.
    /// Returns `None` if the lock to the inner queue cannot be obtained, or if the queue is empty.
    ///
    /// Note: this is IRQ-safe.
    /// TODO: use a data guard to automatically release the lock to the returned InnerQueuedData
    pub fn peek_lockless(&self) -> Option<PeekedData<T>> {
        // first, try to acquire the queue lock.
        let lock_result = self.qref.queue.try_lock();
        match lock_result {
            Some(innerq) => {
                self.peek_locked(innerq)
            },
            None => {
                None
            }
        }
    }

}

/// A producer that can enqueue elements into a DFQueue.
/// Do not wrap this in an Arc or Mutex, the queue it is already protected by those on the interior.
#[derive(Debug)]
pub struct DFQueueProducer<T> {
    qref: Arc<DFQueue<T>>,
    local_backlog: VecDeque<QueuedData<T>>,
}

impl<T> DFQueueProducer<T> {

    /// Call this to obtain another DFQueueProducer.
    /// DFQueueProducer does not implement the standard Clone trait, to avoid accidentally cloning it implicity.
    pub fn obtain_producer(&self) -> DFQueueProducer<T> {
        DFQueueProducer {
            qref: self.qref.clone(),
            local_backlog: VecDeque::new(),
        }
    }

    /// Returns a DFQueueConsumer for this queue, if it hasn't yet been obtained
    /// (either via this function or via `DFQueue::into_consumer()`).
    /// To ensure there is only a single DFQueueConsumer, it will return `None` if there is already a `DFQueueConsumer`.
    pub fn get_consumer(&self) -> Option<DFQueueConsumer<T>> {
        let has_consumer: bool = self.qref.has_consumer.load(Ordering::SeqCst);
        match has_consumer {
            true => None,
            false => {
                self.qref.has_consumer.store(true, Ordering::SeqCst);
                Some(
                    DFQueueConsumer {
                        qref: self.qref.clone(),
                    }
                )
            }
        }
    }

    /// pushes the given `data` onto the back of the queue.
    /// This is a blocking call that will acquire the queue lock, write the data,
    /// and release the lock before returning to the caller.
    ///
    /// # Returns
    /// Returns a QueuedData instance, an Arc-like reference to the given `data` on the queue.
    /// This ensures that the producer can still retain the given `data` if the queue experiences a failure.
    ///
    /// Note: this is NOT IRQ-safe.
    pub fn enqueue_locking(&mut self, data: T) -> QueuedData<T>{

        // acquire the lock so we have exclusive access to the inner queue
        let mut innerq = self.qref.queue.lock(); // blocking call

        // 1) move all the previously-backlogged data into the inner queue
        // 2) push the actual `data` onto the queue
        innerq.append(&mut self.local_backlog);

        let queued_data: QueuedData<T> = QueuedData::new(data);
        innerq.push_back(queued_data.clone());

        // TODO: remove items from the innerq that have been completed

        queued_data
    }

    /// pushes the given `data` onto the back of the queue, only if the queue lock can be acquired immediately.
    /// If the queue lock cannot be acquired, then the given `data` is put on a local backlog queue
    /// that is specific to this `DFQueueProducer`, which will be merged into the actual inner queue
    /// on the next invocation of any `enqueue` method, provided that the inner queue can be locked immediately at that point.
    ///
    /// # Returns
    /// Returns a tuple of the form `(QueuedData, bool)` in which:
    /// <ul> QueuedData is an Arc-like reference to the given `data` on the queue, which
    /// ensures that the producer can still retain the given `data` if the queue experiences a failure.
    /// <ul> bool: true if the given `data` was added to the actual queue,
    ///            false if it was placed into this `DFQueueProducer`'s backlog queue.
    ///
    /// # Important Note
    /// if the returned boolean is false, the data was placed on the backlog queue and
    /// IS NOT GUARANTEED TO BE ON THE ACTUAL QUEUE until [`flush_backlog`]: #method.flush_backlog is called.
    /// Note that calling [`enqueue_locking`]: #method.enqueue_locking will cause the backlog to be flushed to the actual queue.
    /// If `enqueue_lockless` is called again, the backlogged data MAY be added to the actual queue, but that is also not guaranteed!
    ///
    ///
    /// Note: this is IRQ-safe.
    pub fn enqueue_lockless(&mut self, data: T) -> (QueuedData<T>, bool) {

        let queued_data: QueuedData<T> = QueuedData::new(data);

        // first, try to acquire the queue lock.
        match self.qref.queue.try_lock() {
            Some(mut innerq) => {

                // 1) move all the previously-backlogged data into the inner queue
                // 2) push the actual `dataref` onto the queue
                innerq.append(&mut self.local_backlog);
                innerq.push_back(queued_data.clone());

                // TODO: remove completed items from the queue

                (queued_data, true)
            },
            None => {
                // if we can NOT acquire it:
                // add the given `dataref` to this DFQueueRef's producer-local backlog queue

                // println!("enqueue_lockless: backlogging data: {:?}", queued_data);

                self.local_backlog.push_back(queued_data.clone());

                (queued_data, false)
            }
        }
    }

    /// A blocking call that flushes this `DFQueueProducer`'s backlog queue to the actual queue.
    ///
    /// Note: this is NOT IRQ-safe.
    pub fn flush_backlog(&mut self) {
        let mut innerq = self.qref.queue.lock(); // blocking call
        innerq.append(&mut self.local_backlog);
    }

}

// Conditionally compile the module `test` only when the test-suite is run.
#[cfg(test)]
mod test {

    extern crate std;

    use std::thread;
    use super::*;

    #[test]
    // #[should_panic]
    fn simple_test() {

        let queue: DFQueue<usize> = DFQueue::new();

        let mut queue_prod = queue.into_producer();
        let mut queue_cons = queue_prod.get_consumer();
        assert!(queue_cons.is_some(), "First DFQueueConsumer was None!!!");

        // a second call to get_consumer must return None
        assert!(queue_prod.get_consumer().is_none(), "Second DFQueueConsumer wasn't None!!");

        let queue_cons = queue_cons.unwrap();

        let mut queue_prod2 = queue_prod.obtain_producer();

        let mut thr_p = thread::spawn( move || {
            let original_data: Vec<usize> = vec![1, 2, 3, 4, 5];

            for i in 1..20 {
                for elem in original_data.iter() {
                    queue_prod.enqueue_lockless(*elem);
                }
            }

            let (queued_data, backlogged) = queue_prod.enqueue_lockless(256);
            println!("prod1: queued_data = {:?}, backlogged?: {}", queued_data, backlogged);

            queue_prod.flush_backlog();
        } );

         let mut thr_p2 = thread::spawn( move || {
            let original_data: Vec<usize> = vec![10, 11, 12, 13, 14];

            for i in 1..20 {
                for elem in original_data.iter() {
                    queue_prod2.enqueue_lockless(*elem);
                }
            }

            let (queued_data, backlogged) = queue_prod2.enqueue_lockless(512);
            println!("prod2: queued_data = {:?}, backlogged?: {}", queued_data, backlogged);

            queue_prod2.flush_backlog();
        } );

        let mut thr_c = thread::spawn( move || {
            loop {
                let mut val = queue_cons.peek_locking();
                if let Some(v) = val {
                    println!("peeked: {:?}, value={}", v, *v);
                    v.mark_completed();
                }
                else {
                    // println!("\nDumping queue: \n  {:?}", queue_cons.qref.queue);
                }
            }
        } );

        thr_p.join();
        println!("producer 1 done");
        thr_p2.join();
        println!("producer 2 done");
        thr_c.join();
    }
}<|endoftext|>


<commit_before>#[macro_use] extern crate log;

extern crate byteorder;
extern crate dns_parser;
extern crate futures;
extern crate libc;
extern crate multimap;
extern crate net2;
extern crate nix;
extern crate rand;
extern crate tokio_core as tokio;

use dns_parser::Name;
use futures::{BoxFuture, Future};
use futures::sync::mpsc;
use std::io;
use std::thread;
use std::sync::{Arc, RwLock};
use std::cell::RefCell;
use tokio::reactor::{Handle, Core};

mod address_family;
mod fsm;
mod services;
mod net;

use address_family::{Inet, Inet6};
use services::{ServicesInner, Services, ServiceData};
use fsm::{Command, FSM};

const DEFAULT_TTL : u32 = 60;
const MDNS_PORT : u16 = 5353;

pub struct Responder {
    services: Services,
    commands: RefCell<CommandSender>,
    shutdown: Arc<Shutdown>,
}

pub struct Service {
    id: usize,
    services: Services,
    commands: CommandSender,
    _shutdown: Arc<Shutdown>,
}

impl Responder {
    fn setup_core() -> io::Result<(Core, BoxFuture<(), io::Error>, Responder)> {
        let core = Core::new()?;
        let (responder, task) = Self::with_handle(&core.handle())?;
        Ok((core, task, responder))
    }

    pub fn new() -> io::Result<Responder> {
        let (tx, rx) = std::sync::mpsc::sync_channel(0);
        thread::Builder::new()
            .name("mdns-responder".to_owned())
            .spawn(move || {
                match Self::setup_core() {
                    Ok((mut core, task, responder)) => {
                        tx.send(Ok(responder)).expect("tx responder channel closed");
                        core.run(task).expect("mdsn thread failed");
                    }
                    Err(err) => {
                        tx.send(Err(err)).expect("tx responder channel closed");
                    }
                }
            })?;

        rx.recv().expect("rx responder channel closed")
    }

    pub fn spawn(handle: &Handle) -> io::Result<Responder> {
        let (responder, task) = Responder::with_handle(handle)?;
        handle.spawn(task.map_err(|e| {
            warn!("mdns error {:?}", e);
            ()
        }));
        Ok(responder)
    }

    pub fn with_handle(handle: &Handle) -> io::Result<(Responder, BoxFuture<(), io::Error>)> {
        let mut hostname = try!(net::gethostname());
        if !hostname.ends_with(".local") {
            hostname.push_str(".local");
        }

        let services = Arc::new(RwLock::new(ServicesInner::new(hostname)));

        let v4 = FSM::<Inet>::new(handle, &services);
        let v6 = FSM::<Inet6>::new(handle, &services);

        let (task, commands) = match (v4, v6) {
            (Ok((v4_task, v4_command)), Ok((v6_task, v6_command))) => {
                let task = v4_task.join(v6_task).map(|((),())| ()).boxed();
                let commands = vec![v4_command, v6_command];
                (task, commands)
            }

            (Ok((v4_task, v4_command)), Err(err)) => {
                warn!("Failed to register IPv6 receiver: {:?}", err);
                (v4_task.boxed(), vec![v4_command])
            }

            (Err(err), _) => return Err(err),
        };

        let commands = CommandSender(commands);
        let responder = Responder {
            services: services,
            commands: RefCell::new(commands.clone()),
            shutdown: Arc::new(Shutdown(commands)),
        };

        Ok((responder, task.boxed()))
    }
}

impl Responder {
    pub fn register(&self, svc_type: String, svc_name: String, port: u16, txt: &[&str]) -> Service {
        let txt = if txt.is_empty() {
            vec![0]
        } else {
            txt.into_iter().flat_map(|entry| {
                let entry = entry.as_bytes();
                if entry.len() > 255 {
                    panic!("{:?} is too long for a TXT record", entry);
                }
                std::iter::once(entry.len() as u8).chain(entry.into_iter().cloned())
            }).collect()
        };

        let svc = ServiceData {
            typ: Name::from_str(format!("{}.local", svc_type)).unwrap(),
            name: Name::from_str(format!("{}.{}.local", svc_name, svc_type)).unwrap(),
            port: port,
            txt: txt,
        };

        self.commands.borrow_mut()
            .send_unsolicited(svc.clone(), DEFAULT_TTL, true);

        let id = self.services
            .write().unwrap()
            .register(svc);

        Service {
            id: id,
            commands: self.commands.borrow().clone(),
            services: self.services.clone(),
            _shutdown: self.shutdown.clone(),
        }
    }

}

impl Drop for Service {
    fn drop(&mut self) {
        let svc = self.services
            .write().unwrap()
            .unregister(self.id);
        self.commands.send_unsolicited(svc, 0, false);
    }
}

struct Shutdown(CommandSender);
impl Drop for Shutdown {
    fn drop(&mut self) {
        self.0.send_shutdown();
        // TODO wait for tasks to shutdown
    }
}

#[derive(Clone)]
struct CommandSender(Vec<mpsc::UnboundedSender<Command>>);
impl CommandSender {
    fn send(&mut self, cmd: Command) {
        for tx in self.0.iter_mut() {
            tx.send(cmd.clone()).expect("responder died");
        }
    }

    fn send_unsolicited(&mut self, svc: ServiceData, ttl: u32, include_ip: bool) {
        self.send(Command::SendUnsolicited {
            svc: svc,
            ttl: ttl,
            include_ip: include_ip,
        });
    }

    fn send_shutdown(&mut self) {
        self.send(Command::Shutdown);
    }
}
<commit_msg>Fix typo in error message<commit_after>#[macro_use] extern crate log;

extern crate byteorder;
extern crate dns_parser;
extern crate futures;
extern crate libc;
extern crate multimap;
extern crate net2;
extern crate nix;
extern crate rand;
extern crate tokio_core as tokio;

use dns_parser::Name;
use futures::{BoxFuture, Future};
use futures::sync::mpsc;
use std::io;
use std::thread;
use std::sync::{Arc, RwLock};
use std::cell::RefCell;
use tokio::reactor::{Handle, Core};

mod address_family;
mod fsm;
mod services;
mod net;

use address_family::{Inet, Inet6};
use services::{ServicesInner, Services, ServiceData};
use fsm::{Command, FSM};

const DEFAULT_TTL : u32 = 60;
const MDNS_PORT : u16 = 5353;

pub struct Responder {
    services: Services,
    commands: RefCell<CommandSender>,
    shutdown: Arc<Shutdown>,
}

pub struct Service {
    id: usize,
    services: Services,
    commands: CommandSender,
    _shutdown: Arc<Shutdown>,
}

impl Responder {
    fn setup_core() -> io::Result<(Core, BoxFuture<(), io::Error>, Responder)> {
        let core = Core::new()?;
        let (responder, task) = Self::with_handle(&core.handle())?;
        Ok((core, task, responder))
    }

    pub fn new() -> io::Result<Responder> {
        let (tx, rx) = std::sync::mpsc::sync_channel(0);
        thread::Builder::new()
            .name("mdns-responder".to_owned())
            .spawn(move || {
                match Self::setup_core() {
                    Ok((mut core, task, responder)) => {
                        tx.send(Ok(responder)).expect("tx responder channel closed");
                        core.run(task).expect("mdns thread failed");
                    }
                    Err(err) => {
                        tx.send(Err(err)).expect("tx responder channel closed");
                    }
                }
            })?;

        rx.recv().expect("rx responder channel closed")
    }

    pub fn spawn(handle: &Handle) -> io::Result<Responder> {
        let (responder, task) = Responder::with_handle(handle)?;
        handle.spawn(task.map_err(|e| {
            warn!("mdns error {:?}", e);
            ()
        }));
        Ok(responder)
    }

    pub fn with_handle(handle: &Handle) -> io::Result<(Responder, BoxFuture<(), io::Error>)> {
        let mut hostname = try!(net::gethostname());
        if !hostname.ends_with(".local") {
            hostname.push_str(".local");
        }

        let services = Arc::new(RwLock::new(ServicesInner::new(hostname)));

        let v4 = FSM::<Inet>::new(handle, &services);
        let v6 = FSM::<Inet6>::new(handle, &services);

        let (task, commands) = match (v4, v6) {
            (Ok((v4_task, v4_command)), Ok((v6_task, v6_command))) => {
                let task = v4_task.join(v6_task).map(|((),())| ()).boxed();
                let commands = vec![v4_command, v6_command];
                (task, commands)
            }

            (Ok((v4_task, v4_command)), Err(err)) => {
                warn!("Failed to register IPv6 receiver: {:?}", err);
                (v4_task.boxed(), vec![v4_command])
            }

            (Err(err), _) => return Err(err),
        };

        let commands = CommandSender(commands);
        let responder = Responder {
            services: services,
            commands: RefCell::new(commands.clone()),
            shutdown: Arc::new(Shutdown(commands)),
        };

        Ok((responder, task.boxed()))
    }
}

impl Responder {
    pub fn register(&self, svc_type: String, svc_name: String, port: u16, txt: &[&str]) -> Service {
        let txt = if txt.is_empty() {
            vec![0]
        } else {
            txt.into_iter().flat_map(|entry| {
                let entry = entry.as_bytes();
                if entry.len() > 255 {
                    panic!("{:?} is too long for a TXT record", entry);
                }
                std::iter::once(entry.len() as u8).chain(entry.into_iter().cloned())
            }).collect()
        };

        let svc = ServiceData {
            typ: Name::from_str(format!("{}.local", svc_type)).unwrap(),
            name: Name::from_str(format!("{}.{}.local", svc_name, svc_type)).unwrap(),
            port: port,
            txt: txt,
        };

        self.commands.borrow_mut()
            .send_unsolicited(svc.clone(), DEFAULT_TTL, true);

        let id = self.services
            .write().unwrap()
            .register(svc);

        Service {
            id: id,
            commands: self.commands.borrow().clone(),
            services: self.services.clone(),
            _shutdown: self.shutdown.clone(),
        }
    }

}

impl Drop for Service {
    fn drop(&mut self) {
        let svc = self.services
            .write().unwrap()
            .unregister(self.id);
        self.commands.send_unsolicited(svc, 0, false);
    }
}

struct Shutdown(CommandSender);
impl Drop for Shutdown {
    fn drop(&mut self) {
        self.0.send_shutdown();
        // TODO wait for tasks to shutdown
    }
}

#[derive(Clone)]
struct CommandSender(Vec<mpsc::UnboundedSender<Command>>);
impl CommandSender {
    fn send(&mut self, cmd: Command) {
        for tx in self.0.iter_mut() {
            tx.send(cmd.clone()).expect("responder died");
        }
    }

    fn send_unsolicited(&mut self, svc: ServiceData, ttl: u32, include_ip: bool) {
        self.send(Command::SendUnsolicited {
            svc: svc,
            ttl: ttl,
            include_ip: include_ip,
        });
    }

    fn send_shutdown(&mut self) {
        self.send(Command::Shutdown);
    }
}
<|endoftext|>


<commit_before>//! This is a small crate for the [AnyBar](https://github.com/tonsky/AnyBar).
//!
//! # Examples
//! The functioning of this crate is pretty simple. These are the most use cases:
//!
//! ## Using the default port
//! ```
//! # use anybar::*;
//! // create a new AnyBar instance connected to the default port
//! let mut bar = Anybar::default();
//!
//! // set the color
//! bar.set_color(Color::Red).unwrap();
//! ```
//!
//! ## Using a separate port
//! ```
//! # use anybar::*;
//! // Anybar::new() takes the Anybar port as parameter
//! let mut custom_bar = Anybar::new(1708).unwrap();
//! custom_bar.set_color(Color::Exclamation).unwrap();
//! ```
//!
//! ## Additional information
//! ```
//! # use anybar::*;
//! let mut bar = Anybar::default();
//!
//! // after instantiation, the last color is None
//! assert!(bar.color.is_none());
//!
//! bar.set_color(Color::Red).unwrap();
//! // the last color now contains a value
//! // assert_eq!(bar.color.unwrap(), Color::Red);
//! ```
//!
//! Note that the value of `bar.color` does not necessarily represent the real color
//! displayed at the moment, depending on whether you have something or someone else messing
//! with your AnyBar simultaneously.
//!
//! # Note
//! The AnyBar itself does not provide any information on whether the sent command was executed
//! successfully, so the lib will only return an error if it was not able to bind a UDP socket.
//!

use std::net;

/// The Anybar handle.
pub struct Anybar {
    /// The UDP Port the Anybar is connected to
    pub port: u16,
    /// The last color that has been set.
    ///
    /// When no color has been set yet, color is `None`.
    pub color: Option<Color>,
}

/// The different colors supported by AnyBar.
pub enum Color {
    /// White dot
    White,
    /// Red dot
    Red,
    /// Orange dot
    Orange,
    /// Yellow dot
    Yellow,
    /// Green dot
    Green,
    /// Cyan dot
    Cyan,
    /// Blue dot
    Blue,
    /// Purple dot
    Purple,
    /// Black dot; Has a white frame in _dark_ mode
    Black,
    /// Question mark
    Question,
    /// White exclamation mark on red ground
    Exclamation,
}

impl Anybar {
    /// Create a new Anybar instance, connected to the given UDP port.
    ///
    /// `port` may be any port between 0 and 6553.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anybar::*;
    /// let custom_bar = Anybar::new(1708).unwrap();
    ///
    /// assert_eq!(custom_bar.port, 1708);
    /// ```
    pub fn new(port: u16) -> Result<Anybar, String> {
        if port > 6553 {
            Err(format!("The port {} is not between 0 and 6553!", port))
        } else {
            Ok(Anybar{port:port, color:None})
        }
    }

    fn parse_color(color: &Color) -> &'static [u8] {
        use Color::*;
        match *color {
            White       => b"white",
            Red         => b"red",
            Orange      => b"orange",
            Yellow      => b"yellow",
            Green       => b"green",
            Cyan        => b"cyan",
            Blue        => b"blue",
            Purple      => b"purple",
            Black       => b"black",
            Question    => b"question",
            Exclamation => b"exclamation",
        }
    }

    fn socket(ip: &str, port: u16) -> Result<net::UdpSocket, std::io::Error> {
        net::UdpSocket::bind((ip, port))
    }

    /// Set a new color.
    ///
    /// The function returns a `ResultType` which will contain a `std::io::Error` when the UDP socket can't be bound.
    pub fn set_color(&mut self, color: Color) -> Result<(), std::io::Error> {
        let message = Anybar::parse_color(&color);

        Self::socket("127.0.0.1", 0)?
            .send_to(&message, ("127.0.0.1", self.port))?;

        self.color = Some(color);
        Ok(())
    }

    /// Sends the quit signal to the Anybar and takes ownership of the object.
    ///
    /// _This is an experimental feature of AnyBar._ Since the AnyBar will quit during the
    /// execution of this function, it takes the ownership of `self`,
    /// which will be dropped when this function returns.
    ///
    /// The function returns a `ResultType` which will contain a `std::io::Error` when the UDP socket can't be bound.
    ///
    /// # Example
    /// ```ignore
    /// # use anybar::*;
    /// let mut bar = Anybar::default();
    ///
    /// // do stuff...
    ///
    /// bar.quit().unwrap();
    ///
    /// bar.set_color(Color::White).unwrap();  // this won't work, bar has been moved
    /// ```
    pub fn quit(self) -> Result<(), std::io::Error> {
        Self::socket("127.0.0.1", 0)?
            .send_to(b"quit", ("127.0.0.1", self.port))?;
        Ok(())
    }
}

/// Instanciates the default AnyBar, connected to the Port `1738`.
///
/// ```
/// # use anybar::*;
/// let bar = Anybar::default();
/// assert_eq!(bar.port, 1738);
/// ```
impl Default for Anybar {
    fn default() -> Self {
        Self{port:1738, color:None}
    }
}
<commit_msg>Builds now on stable, removed Self{}<commit_after>//! This is a small crate for the [AnyBar](https://github.com/tonsky/AnyBar).
//!
//! # Examples
//! The functioning of this crate is pretty simple. These are the most use cases:
//!
//! ## Using the default port
//! ```
//! # use anybar::*;
//! // create a new AnyBar instance connected to the default port
//! let mut bar = Anybar::default();
//!
//! // set the color
//! bar.set_color(Color::Red).unwrap();
//! ```
//!
//! ## Using a separate port
//! ```
//! # use anybar::*;
//! // Anybar::new() takes the Anybar port as parameter
//! let mut custom_bar = Anybar::new(1708).unwrap();
//! custom_bar.set_color(Color::Exclamation).unwrap();
//! ```
//!
//! ## Additional information
//! ```
//! # use anybar::*;
//! let mut bar = Anybar::default();
//!
//! // after instantiation, the last color is None
//! assert!(bar.color.is_none());
//!
//! bar.set_color(Color::Red).unwrap();
//! // the last color now contains a value
//! // assert_eq!(bar.color.unwrap(), Color::Red);
//! ```
//!
//! Note that the value of `bar.color` does not necessarily represent the real color
//! displayed at the moment, depending on whether you have something or someone else messing
//! with your AnyBar simultaneously.
//!
//! # Note
//! The AnyBar itself does not provide any information on whether the sent command was executed
//! successfully, so the lib will only return an error if it was not able to bind a UDP socket.
//!

use std::net;

/// The Anybar handle.
pub struct Anybar {
    /// The UDP Port the Anybar is connected to
    pub port: u16,
    /// The last color that has been set.
    ///
    /// When no color has been set yet, color is `None`.
    pub color: Option<Color>,
}

/// The different colors supported by AnyBar.
pub enum Color {
    /// White dot
    White,
    /// Red dot
    Red,
    /// Orange dot
    Orange,
    /// Yellow dot
    Yellow,
    /// Green dot
    Green,
    /// Cyan dot
    Cyan,
    /// Blue dot
    Blue,
    /// Purple dot
    Purple,
    /// Black dot; Has a white frame in _dark_ mode
    Black,
    /// Question mark
    Question,
    /// White exclamation mark on red ground
    Exclamation,
}

impl Anybar {
    /// Create a new Anybar instance, connected to the given UDP port.
    ///
    /// `port` may be any port between 0 and 6553.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anybar::*;
    /// let custom_bar = Anybar::new(1708).unwrap();
    ///
    /// assert_eq!(custom_bar.port, 1708);
    /// ```
    pub fn new(port: u16) -> Result<Anybar, String> {
        if port > 6553 {
            Err(format!("The port {} is not between 0 and 6553!", port))
        } else {
            Ok(Anybar{port:port, color:None})
        }
    }

    fn parse_color(color: &Color) -> &'static [u8] {
        use Color::*;
        match *color {
            White       => b"white",
            Red         => b"red",
            Orange      => b"orange",
            Yellow      => b"yellow",
            Green       => b"green",
            Cyan        => b"cyan",
            Blue        => b"blue",
            Purple      => b"purple",
            Black       => b"black",
            Question    => b"question",
            Exclamation => b"exclamation",
        }
    }

    fn socket(ip: &str, port: u16) -> Result<net::UdpSocket, std::io::Error> {
        net::UdpSocket::bind((ip, port))
    }

    /// Set a new color.
    ///
    /// The function returns a `ResultType` which will contain a `std::io::Error` when the UDP socket can't be bound.
    pub fn set_color(&mut self, color: Color) -> Result<(), std::io::Error> {
        let message = Anybar::parse_color(&color);

        Self::socket("127.0.0.1", 0)?
            .send_to(&message, ("127.0.0.1", self.port))?;

        self.color = Some(color);
        Ok(())
    }

    /// Sends the quit signal to the Anybar and takes ownership of the object.
    ///
    /// _This is an experimental feature of AnyBar._ Since the AnyBar will quit during the
    /// execution of this function, it takes the ownership of `self`,
    /// which will be dropped when this function returns.
    ///
    /// The function returns a `ResultType` which will contain a `std::io::Error` when the UDP socket can't be bound.
    ///
    /// # Example
    /// ```ignore
    /// # use anybar::*;
    /// let mut bar = Anybar::default();
    ///
    /// // do stuff...
    ///
    /// bar.quit().unwrap();
    ///
    /// bar.set_color(Color::White).unwrap();  // this won't work, bar has been moved
    /// ```
    pub fn quit(self) -> Result<(), std::io::Error> {
        Self::socket("127.0.0.1", 0)?
            .send_to(b"quit", ("127.0.0.1", self.port))?;
        Ok(())
    }
}

/// Instanciates the default AnyBar, connected to the Port `1738`.
///
/// ```
/// # use anybar::*;
/// let bar = Anybar::default();
/// assert_eq!(bar.port, 1738);
/// ```
impl Default for Anybar {
    fn default() -> Self {
        Anybar{port:1738, color:None}
    }
}
<|endoftext|>


<commit_before>extern crate log;
extern crate simplelog;

use goblin::pe::PE;
use goblin::Object;
use log::{debug, error, info, trace};
use rayon::prelude::*;
use std::env;
use std::fs;
use std::io::prelude::*;
use zydis;

// TODO: this really shouldn't be public. only for testing.
// TODO: what is the `cfg(doctest)` setting?
//#[cfg(test)]
pub mod rsrc;

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

#[derive(Debug)]
pub enum Error {
    FileAccess,
    FileFormat,
    NotImplemented,
    ValueError,
}

/// Round the given value up to the next multiple of the given base.
///
/// # Panics
///
///   - Base `b` must be at least `2`.
///
/// # Examples
///
/// ```
/// use lancelot::*;
/// assert_eq!(align(0, 2), 0);
/// assert_eq!(align(1, 2), 2);
/// assert_eq!(align(2, 2), 2);
/// assert_eq!(align(3, 2), 4);
/// assert_eq!(align(4, 2), 4);
/// ```
pub fn align(i: usize, b: usize) -> usize {
    if b < 2 {
        panic!("base `b` must be at least: 2");
    }
    let rem = i % b;
    if rem == 0 {
        i
    } else {
        i + (b - rem)
    }
}

pub fn hexdump_ascii(b: u8) -> char {
    if b.is_ascii_graphic() || b == b' ' {
        b as char
    } else {
        '.'
    }
}

pub fn hexdump(buf: &[u8], offset: usize) -> String {
    // 01234567:  00 01 02 03 04 05 06 07  ...............
    // <prefix>   <hex col>                <ascii col>

    let padding = "  ";

    let padding_size = 2;
    let hex_col_size = 3;
    let ascii_col_size = 1;
    let prefix_size = 8 + 1;
    let newline_size = 1;
    let line_size = prefix_size
        + padding_size
        + 16 * hex_col_size
        + padding_size
        + 16 * ascii_col_size
        + newline_size;
    let line_count = align(buf.len(), 0x10) / 0x10;

    let mut ret = String::with_capacity(line_count * line_size);

    let mut line = String::with_capacity(line_size);
    let mut remaining_count = buf.len();
    for line_index in 0..line_count {
        let line_elem_count = 0x10.min(remaining_count);
        let padding_elem_count = 0x10 - line_elem_count;

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        // ^^^^^^^^^
        line.push_str(format!("{:08x}:", offset + 0x10 * line_index).as_str());

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //          ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //            ^^^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push_str(format!("{:02x} ", elem).as_str());
        }
        for _ in 0..padding_elem_count {
            line.push_str("   ");
        }

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                   ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                     ^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push(hexdump_ascii(*elem))
        }
        for _ in 0..padding_elem_count {
            line.push(' ');
        }
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                                    ^
        line.push('\n');

        ret.push_str(line.as_str());
        line.truncate(0x0);
        remaining_count -= line_elem_count;
    }

    ret
}

fn foo(ws: &Workspace) -> Result<(), Error> {
    if let Object::PE(pe) = ws.get_obj().unwrap() {
        info!("module: {}", pe.name.unwrap_or("(unknown)"));

        info!("bitness: {}", if pe.is_64 { "64" } else { "32" });
        info!("image base: 0x{:x}", pe.image_base);
        info!("entry rva: 0x{:x}", pe.entry);

        // like:
        //
        //     sections:
        //       - .text
        //         raw size:     0x18aa00
        //         virtual size: 0x18a9a8
        info!("sections:");
        for section in ws.sections.iter() {
            info!("  - {}", section.name);
            info!("    virtual address: 0x{:x}", section.addr);
            info!("    virtual size: 0x{:x}", section.buf.len());

            info!(
                "\n{}",
                hexdump(&section.buf[..0x1C], pe.image_base + section.addr as usize)
            );

            let decoder =
                zydis::Decoder::new(zydis::MachineMode::Long64, zydis::AddressWidth::_64).unwrap();
            let insns: Vec<_> = section
                .buf
                .par_windows(0x10)
                .map(|ibuf| decoder.decode(ibuf))
                .collect();

            info!("total instructions: {}", insns.len());

            info!(
                "successful disassembles: {}",
                insns.par_iter().filter(|insn| insn.is_ok()).count()
            );

            info!(
                "valid instructions: {}",
                insns
                    .par_iter()
                    .filter(|insn| match insn {
                        Ok(Some(_)) => true,
                        _ => false,
                    })
                    .count()
            );
        }
    }

    Ok(())
}

pub fn read_file(filename: &str) -> Result<Vec<u8>, Error> {
    debug!("read_file: {:?}", filename);

    let mut buf = Vec::new();
    {
        debug!("reading file: {}", filename);
        let mut f = match fs::File::open(filename) {
            Ok(f) => f,
            Err(_) => {
                error!("failed to open file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        let bytes_read = match f.read_to_end(&mut buf) {
            Ok(c) => c,
            Err(_) => {
                error!("failed to read entire file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        debug!("read {} bytes", bytes_read);
        if bytes_read < 0x10 {
            error!("file too small: {}", filename);
            return Err(Error::FileFormat);
        }
    }

    Ok(buf)
}

pub struct Section {
    pub name: String,
    pub addr: u32,
    pub buf: Vec<u8>,
    pub insns: Vec<Option<zydis::ffi::DecodedInstruction>>,
}

pub struct Workspace {
    pub filename: String,
    pub buf: Vec<u8>,
    pub sections: Vec<Section>,
}

impl Workspace {
    /// Parse the given file into its object.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///
    /// # Examples
    ///
    /// ```
    /// use goblin::Object;
    /// use matches::matches;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// assert!(matches!(ws.get_obj().unwrap(), Object::PE(_)));
    /// ```
    ///
    /// you might be tempted to maintain a method `get_pe`,
    /// however, i don't think this is a good idea:
    /// its fragile because the file type may not be PE.
    /// therefore, force clients to be explicit:
    ///
    /// ```
    /// use goblin::Object;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// if let Object::PE(_) = ws.get_obj().unwrap() {
    ///     // everyone is happy!
    /// }
    /// ```
    ///
    /// TODO: demonstrate `Error::FileFormat`.
    pub fn get_obj(&self) -> Result<Object, Error> {
        let obj = match Object::parse(&self.buf) {
            Ok(o) => o,
            Err(e) => {
                error!("failed to parse file: {} error: {:?}", self.filename, e);
                return Err(Error::FileFormat);
            }
        };

        match obj {
            Object::Unknown(_) => {
                error!(
                    "unknown file format, magic: | {:02X} {:02X} | '{}{}' ",
                    self.buf[0],
                    self.buf[1],
                    hexdump_ascii(self.buf[0]),
                    hexdump_ascii(self.buf[1])
                );

                Err(Error::FileFormat)
            }
            _ => Ok(obj),
        }
    }

    pub fn get_pe(&self) -> Result<PE, Error> {
        match self.get_obj()? {
            Object::PE(pe) => {
                return Ok(pe);
            }
            _ => {
                error!("not a PE file");
                return Err(Error::ValueError);
            }
        }
    }

    /// Construct a workspace from the module with the given contents.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///   - `Error::NotImplemented`: when not a PE file.
    ///
    /// # Examples
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let buf = get_buf(Rsrc::K32);
    /// let ws = Workspace::from_buf("kernel32.dll", buf).unwrap();
    /// ```
    ///
    /// TODO: demonstrate ELF file behavior.
    /// TODO: demonstrate MachO file behavior.
    /// TODO: demonstrate unknown file behavior.
    pub fn from_buf(filename: &str, buf: Vec<u8>) -> Result<Workspace, Error> {
        let mut ws = Workspace {
            filename: filename.to_string(),
            buf: buf.clone(),
            sections: vec![],
        };

        match ws.get_obj()? {
            Object::PE(pe) => {
                info!("found PE file");

                let machine;
                let mode;
                if pe.is_64 {
                    machine = zydis::MachineMode::Long64;
                    mode = zydis::AddressWidth::_64;
                } else {
                    // TODO: not sure what `LongCompat32` means, vs `Legacy32`.
                    machine = zydis::MachineMode::LongCompat32;
                    mode = zydis::AddressWidth::_32;
                }
                // TODO: save off decoder into workspace.
                let decoder = zydis::Decoder::new(machine, mode).unwrap();

                // TODO: load PE header, too

                ws.sections
                    .extend(pe.sections.iter().map(|section| -> Section {
                        // TODO: i'm sure this can be abused.
                        // TODO: add tests for weird section names.
                        let name = String::from_utf8_lossy(&section.name[..]).into_owned();

                        // TODO: figure out if we will work with usize, or u64, or what,
                        // then assert usize is ok.
                        // ref: `usize::max_value()`
                        let mut secbuf = vec![0; align(section.virtual_size as usize, 0x200)];

                        {
                            let secsize = section.size_of_raw_data as usize;
                            let rawbuf = &mut secbuf[..secsize];
                            let pstart = section.pointer_to_raw_data as usize;
                            rawbuf.copy_from_slice(&buf[pstart..pstart + secsize]);
                        }

                        let insns: Vec<_> = secbuf
                            .par_windows(0x10)
                            .map(|ibuf| match decoder.decode(ibuf) {
                                Ok(Some(insn)) => Some(insn),
                                _ => None,
                            })
                            .collect();

                        Section {
                            name: name,
                            addr: section.virtual_address,
                            buf: secbuf,
                            insns: insns,
                        }
                    }));
            }
            Object::Elf(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Mach(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Archive(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Unknown(_) => {
                return Err(Error::NotImplemented);
            }
        }

        Ok(ws)
    }

    /// Construct a workspace from the module at given a file path.
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let path = get_path(Rsrc::K32);
    /// // This test resource file is mangled. Needs to be fixed before parsing.
    /// // Otherwise, the following would work:
    /// // let ws = Workspace::from_file(&path).unwrap();
    /// ```
    pub fn from_file(filename: &str) -> Result<Workspace, Error> {
        let buf = read_file(filename)?;
        Workspace::from_buf(filename, buf)
    }
}

pub fn run(args: &Config) -> Result<(), Error> {
    debug!("filename: {:?}", args.filename);
    let ws = Workspace::from_file(&args.filename)?;

    if let Object::PE(_) = ws.get_obj()? {
        foo(&ws).expect("failed to foo");
    }

    Ok(())
}
<commit_msg>lib: parse section name more correctly<commit_after>extern crate log;
extern crate simplelog;

use goblin::pe::PE;
use goblin::Object;
use log::{debug, error, info, trace};
use rayon::prelude::*;
use std::env;
use std::fs;
use std::io::prelude::*;
use zydis;

// TODO: this really shouldn't be public. only for testing.
// TODO: what is the `cfg(doctest)` setting?
//#[cfg(test)]
pub mod rsrc;

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

#[derive(Debug)]
pub enum Error {
    FileAccess,
    FileFormat,
    NotImplemented,
    ValueError,
}

/// Round the given value up to the next multiple of the given base.
///
/// # Panics
///
///   - Base `b` must be at least `2`.
///
/// # Examples
///
/// ```
/// use lancelot::*;
/// assert_eq!(align(0, 2), 0);
/// assert_eq!(align(1, 2), 2);
/// assert_eq!(align(2, 2), 2);
/// assert_eq!(align(3, 2), 4);
/// assert_eq!(align(4, 2), 4);
/// ```
pub fn align(i: usize, b: usize) -> usize {
    if b < 2 {
        panic!("base `b` must be at least: 2");
    }
    let rem = i % b;
    if rem == 0 {
        i
    } else {
        i + (b - rem)
    }
}

pub fn hexdump_ascii(b: u8) -> char {
    if b.is_ascii_graphic() || b == b' ' {
        b as char
    } else {
        '.'
    }
}

pub fn hexdump(buf: &[u8], offset: usize) -> String {
    // 01234567:  00 01 02 03 04 05 06 07  ...............
    // <prefix>   <hex col>                <ascii col>

    let padding = "  ";

    let padding_size = 2;
    let hex_col_size = 3;
    let ascii_col_size = 1;
    let prefix_size = 8 + 1;
    let newline_size = 1;
    let line_size = prefix_size
        + padding_size
        + 16 * hex_col_size
        + padding_size
        + 16 * ascii_col_size
        + newline_size;
    let line_count = align(buf.len(), 0x10) / 0x10;

    let mut ret = String::with_capacity(line_count * line_size);

    let mut line = String::with_capacity(line_size);
    let mut remaining_count = buf.len();
    for line_index in 0..line_count {
        let line_elem_count = 0x10.min(remaining_count);
        let padding_elem_count = 0x10 - line_elem_count;

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        // ^^^^^^^^^
        line.push_str(format!("{:08x}:", offset + 0x10 * line_index).as_str());

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //          ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //            ^^^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push_str(format!("{:02x} ", elem).as_str());
        }
        for _ in 0..padding_elem_count {
            line.push_str("   ");
        }

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                   ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                     ^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push(hexdump_ascii(*elem))
        }
        for _ in 0..padding_elem_count {
            line.push(' ');
        }
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                                    ^
        line.push('\n');

        ret.push_str(line.as_str());
        line.truncate(0x0);
        remaining_count -= line_elem_count;
    }

    ret
}

fn foo(ws: &Workspace) -> Result<(), Error> {
    if let Object::PE(pe) = ws.get_obj().unwrap() {
        info!("module: {}", pe.name.unwrap_or("(unknown)"));

        info!("bitness: {}", if pe.is_64 { "64" } else { "32" });
        info!("image base: 0x{:x}", pe.image_base);
        info!("entry rva: 0x{:x}", pe.entry);

        // like:
        //
        //     sections:
        //       - .text
        //         raw size:     0x18aa00
        //         virtual size: 0x18a9a8
        info!("sections:");
        for section in ws.sections.iter() {
            info!("  - {}", section.name);
            info!("    virtual address: 0x{:x}", section.addr);
            info!("    virtual size: 0x{:x}", section.buf.len());

            info!(
                "\n{}",
                hexdump(&section.buf[..0x1C], pe.image_base + section.addr as usize)
            );

            let decoder =
                zydis::Decoder::new(zydis::MachineMode::Long64, zydis::AddressWidth::_64).unwrap();
            let insns: Vec<_> = section
                .buf
                .par_windows(0x10)
                .map(|ibuf| decoder.decode(ibuf))
                .collect();

            info!("total instructions: {}", insns.len());

            info!(
                "successful disassembles: {}",
                insns.par_iter().filter(|insn| insn.is_ok()).count()
            );

            info!(
                "valid instructions: {}",
                insns
                    .par_iter()
                    .filter(|insn| match insn {
                        Ok(Some(_)) => true,
                        _ => false,
                    })
                    .count()
            );
        }
    }

    Ok(())
}

pub fn read_file(filename: &str) -> Result<Vec<u8>, Error> {
    debug!("read_file: {:?}", filename);

    let mut buf = Vec::new();
    {
        debug!("reading file: {}", filename);
        let mut f = match fs::File::open(filename) {
            Ok(f) => f,
            Err(_) => {
                error!("failed to open file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        let bytes_read = match f.read_to_end(&mut buf) {
            Ok(c) => c,
            Err(_) => {
                error!("failed to read entire file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        debug!("read {} bytes", bytes_read);
        if bytes_read < 0x10 {
            error!("file too small: {}", filename);
            return Err(Error::FileFormat);
        }
    }

    Ok(buf)
}

pub struct Section {
    pub name: String,
    pub addr: u32,
    pub buf: Vec<u8>,
    pub insns: Vec<Option<zydis::ffi::DecodedInstruction>>,
}

pub struct Workspace {
    pub filename: String,
    pub buf: Vec<u8>,
    pub sections: Vec<Section>,
}

impl Workspace {
    /// Parse the given file into its object.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///
    /// # Examples
    ///
    /// ```
    /// use goblin::Object;
    /// use matches::matches;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// assert!(matches!(ws.get_obj().unwrap(), Object::PE(_)));
    /// ```
    ///
    /// you might be tempted to maintain a method `get_pe`,
    /// however, i don't think this is a good idea:
    /// its fragile because the file type may not be PE.
    /// therefore, force clients to be explicit:
    ///
    /// ```
    /// use goblin::Object;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// if let Object::PE(_) = ws.get_obj().unwrap() {
    ///     // everyone is happy!
    /// }
    /// ```
    ///
    /// TODO: demonstrate `Error::FileFormat`.
    pub fn get_obj(&self) -> Result<Object, Error> {
        let obj = match Object::parse(&self.buf) {
            Ok(o) => o,
            Err(e) => {
                error!("failed to parse file: {} error: {:?}", self.filename, e);
                return Err(Error::FileFormat);
            }
        };

        match obj {
            Object::Unknown(_) => {
                error!(
                    "unknown file format, magic: | {:02X} {:02X} | '{}{}' ",
                    self.buf[0],
                    self.buf[1],
                    hexdump_ascii(self.buf[0]),
                    hexdump_ascii(self.buf[1])
                );

                Err(Error::FileFormat)
            }
            _ => Ok(obj),
        }
    }

    pub fn get_pe(&self) -> Result<PE, Error> {
        match self.get_obj()? {
            Object::PE(pe) => {
                return Ok(pe);
            }
            _ => {
                error!("not a PE file");
                return Err(Error::ValueError);
            }
        }
    }

    /// Construct a workspace from the module with the given contents.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///   - `Error::NotImplemented`: when not a PE file.
    ///
    /// # Examples
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let buf = get_buf(Rsrc::K32);
    /// let ws = Workspace::from_buf("kernel32.dll", buf).unwrap();
    /// ```
    ///
    /// TODO: demonstrate ELF file behavior.
    /// TODO: demonstrate MachO file behavior.
    /// TODO: demonstrate unknown file behavior.
    pub fn from_buf(filename: &str, buf: Vec<u8>) -> Result<Workspace, Error> {
        let mut ws = Workspace {
            filename: filename.to_string(),
            buf: buf.clone(),
            sections: vec![],
        };

        match ws.get_obj()? {
            Object::PE(pe) => {
                info!("found PE file");

                let machine;
                let mode;
                if pe.is_64 {
                    machine = zydis::MachineMode::Long64;
                    mode = zydis::AddressWidth::_64;
                } else {
                    // TODO: not sure what `LongCompat32` means, vs `Legacy32`.
                    machine = zydis::MachineMode::LongCompat32;
                    mode = zydis::AddressWidth::_32;
                }
                // TODO: save off decoder into workspace.
                let decoder = zydis::Decoder::new(machine, mode).unwrap();

                // TODO: load PE header, too

                ws.sections
                    .extend(pe.sections.iter().map(|section| -> Section {
                        // TODO: i'm sure this can be abused.
                        // TODO: add tests for weird section names.
                        let name = String::from_utf8_lossy(&section.name[..])
                            .into_owned()
                            .trim_end_matches("\u{0}")
                            .trim_end()
                            .to_string();

                        // TODO: figure out if we will work with usize, or u64, or what,
                        // then assert usize is ok.
                        // ref: `usize::max_value()`
                        let mut secbuf = vec![0; align(section.virtual_size as usize, 0x200)];

                        {
                            let secsize = section.size_of_raw_data as usize;
                            let rawbuf = &mut secbuf[..secsize];
                            let pstart = section.pointer_to_raw_data as usize;
                            rawbuf.copy_from_slice(&buf[pstart..pstart + secsize]);
                        }

                        let insns: Vec<_> = secbuf
                            .par_windows(0x10)
                            .map(|ibuf| match decoder.decode(ibuf) {
                                Ok(Some(insn)) => Some(insn),
                                _ => None,
                            })
                            .collect();

                        Section {
                            name: name,
                            addr: section.virtual_address,
                            buf: secbuf,
                            insns: insns,
                        }
                    }));
            }
            Object::Elf(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Mach(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Archive(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Unknown(_) => {
                return Err(Error::NotImplemented);
            }
        }

        Ok(ws)
    }

    /// Construct a workspace from the module at given a file path.
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let path = get_path(Rsrc::K32);
    /// // This test resource file is mangled. Needs to be fixed before parsing.
    /// // Otherwise, the following would work:
    /// // let ws = Workspace::from_file(&path).unwrap();
    /// ```
    pub fn from_file(filename: &str) -> Result<Workspace, Error> {
        let buf = read_file(filename)?;
        Workspace::from_buf(filename, buf)
    }
}

pub fn run(args: &Config) -> Result<(), Error> {
    debug!("filename: {:?}", args.filename);
    let ws = Workspace::from_file(&args.filename)?;

    if let Object::PE(_) = ws.get_obj()? {
        foo(&ws).expect("failed to foo");
    }

    Ok(())
}
<|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,
    Warn {message: String},
    Abort {message: String},
}

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

#[derive(Debug)]
enum Step {
    Echo {name: String, message: String},
    Run {name: String, command: String, arguments: String, on_success: OnSuccess, on_failure: OnFailure},
    Shell,
    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
    // unimplemented!();

    // 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() {
                            "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() {
                            "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" => {
            // },
            _ => {
                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();
                let arguments: Vec<&str> = arguments.split(' ').collect();
                println!("    command: {}   arguments: {:?}", command, arguments);
                let mut 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() {
                    println!("Success!");
                    match on_success {
                        OnSuccess::Continue => {},
                        OnSuccess::Warn {message} => {
                            println!("WARNING: {}", message);
                        },
                        OnSuccess::Abort {message} => {
                            println!("ABORT: {}", message);
                            break;
                        },
                    }
                } else {
                    match on_failure {
                        OnFailure::Continue => {},
                        OnFailure::Warn {message} => {
                            println!("WARNING: {}", message);
                        },
                        OnFailure::Abort {message} => {
                            println!("ABORT: {}", message);
                            break;
                        },
                    }
                }
            },
            Step::SetEnv {variable, value} => {
                std::env::set_var(variable, value);
            },
            _ => unimplemented!(),
        }
    }

}
<commit_msg>Set environment variables from 'variables' table<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,
    Warn {message: String},
    Abort {message: String},
}

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

#[derive(Debug)]
enum Step {
    Echo {name: String, message: String},
    Run {name: String, command: String, arguments: String, on_success: OnSuccess, on_failure: OnFailure},
    Shell,
    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() {
                            "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() {
                            "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" => {
            // },
            _ => {
                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();
                let arguments: Vec<&str> = arguments.split(' ').collect();
                println!("    command: {}   arguments: {:?}", command, arguments);
                let mut 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() {
                    println!("Success!");
                    match on_success {
                        OnSuccess::Continue => {},
                        OnSuccess::Warn {message} => {
                            println!("WARNING: {}", message);
                        },
                        OnSuccess::Abort {message} => {
                            println!("ABORT: {}", message);
                            break;
                        },
                    }
                } else {
                    match on_failure {
                        OnFailure::Continue => {},
                        OnFailure::Warn {message} => {
                            println!("WARNING: {}", message);
                        },
                        OnFailure::Abort {message} => {
                            println!("ABORT: {}", message);
                            break;
                        },
                    }
                }
            },
            Step::SetEnv {variable, value} => {
                std::env::set_var(variable, value);
            },
            _ => unimplemented!(),
        }
    }

}
<|endoftext|>


<commit_before>extern crate libc;

use std::io;

mod ffi;

macro_rules! unsafe_try {
    ( $x:expr ) => {
        try!(int_result(unsafe { $x }))
    };
}

pub fn fork() -> io::Result<(libc::pid_t, libc::c_int)>
{
    let pty_master = try!(open_ptm());
    let pid        = unsafe_try!(libc::fork());

    if pid == 0 {
        try!(attach_pts(pty_master));

        return Ok((0, -1));
    }
    else {
        return Ok((pid, pty_master));
    }
}

fn open_ptm() -> io::Result<libc::c_int> {
    let pty_master = unsafe_try!(ffi::posix_openpt(libc::O_RDWR));

    unsafe_try!(ffi::grantpt(pty_master));
    unsafe_try!(ffi::unlockpt(pty_master));

    Ok(pty_master)
}

fn attach_pts(pty_master: libc::c_int) -> io::Result<()> {
    let pts_name = unsafe { ffi::ptsname(pty_master) };

    if (pts_name as i32) == 0 {
        return Err(io::Error::last_os_error())
    }

    unsafe_try!(libc::close(pty_master));
    unsafe_try!(libc::setsid());

    let pty_slave = unsafe_try!(libc::open(pts_name, libc::O_RDWR, 0));

    unsafe_try!(libc::dup2(pty_slave, libc::STDIN_FILENO));
    unsafe_try!(libc::dup2(pty_slave, libc::STDOUT_FILENO));
    unsafe_try!(libc::dup2(pty_slave, libc::STDERR_FILENO));

    unsafe_try!(libc::close(pty_slave));

    Ok(())
}

#[inline]
fn int_result(value: libc::c_int) -> io::Result<libc::c_int> {
    if value < 0 {
        return Err(io::Error::last_os_error())
    }

    Ok(value)
}

#[test]
fn it_works() {
}
<commit_msg>Return an object, not a raw value<commit_after>#![feature(zero_one)]

extern crate libc;

use std::io;
use std::num::One;
use std::ops::Neg;

mod ffi;

macro_rules! unsafe_try {
    ( $x:expr ) => {
        try!(to_result(unsafe { $x }))
    };
}

pub struct Child {
    pid: libc::pid_t
}

impl Child {
    pub fn pid(&self) -> libc::pid_t {
        self.pid
    }

    pub fn wait(&self) {
        unsafe { libc::waitpid(self.pid, std::ptr::null(), 0) };
    }
}

pub struct Master {
    fd: libc::c_int
}

impl Master {
    pub fn raw(&self) -> libc::c_int {
        self.fd
    }

    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
        let ret = unsafe_try!(
            libc::read(self.fd,
                       buf.as_mut_ptr() as *mut libc::c_void,
                       buf.len() as libc::size_t)
        );

        Ok(ret as usize)
    }

    pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
        let ret = unsafe_try!(
            libc::write(self.fd,
                        buf.as_ptr() as *const libc::c_void,
                        buf.len() as libc::size_t)
        );

        Ok(ret as usize)
    }

    pub fn close(&self) {
        unsafe { libc::close(self.fd) };
    }
}

impl Copy for Master {}

impl Clone for Master {
    fn clone(&self) -> Master { *self }
}

pub fn fork() -> io::Result<(Child, Master)>
{
    let pty_master = try!(open_ptm());
    let pid        = unsafe_try!(libc::fork());

    if pid == 0 {
        try!(attach_pts(pty_master));

        return Ok((Child { pid: 0 }, Master { fd: -1 }));
    }
    else {
        return Ok((Child { pid: pid }, Master { fd: pty_master }));
    }
}

fn open_ptm() -> io::Result<libc::c_int> {
    let pty_master = unsafe_try!(ffi::posix_openpt(libc::O_RDWR));

    unsafe_try!(ffi::grantpt(pty_master));
    unsafe_try!(ffi::unlockpt(pty_master));

    Ok(pty_master)
}

fn attach_pts(pty_master: libc::c_int) -> io::Result<()> {
    let pts_name = unsafe { ffi::ptsname(pty_master) };

    if (pts_name as i32) == 0 {
        return Err(io::Error::last_os_error())
    }

    unsafe_try!(libc::close(pty_master));
    unsafe_try!(libc::setsid());

    let pty_slave = unsafe_try!(libc::open(pts_name, libc::O_RDWR, 0));

    unsafe_try!(libc::dup2(pty_slave, libc::STDIN_FILENO));
    unsafe_try!(libc::dup2(pty_slave, libc::STDOUT_FILENO));
    unsafe_try!(libc::dup2(pty_slave, libc::STDERR_FILENO));

    unsafe_try!(libc::close(pty_slave));

    Ok(())
}

#[inline]
fn to_result<T: One + PartialEq + Neg<Output=T>>(t: T) -> io::Result<T> {
    let one: T = T::one();

    if t == -one {
        Err(io::Error::last_os_error())
    } else {
        Ok(t)
    }
}

#[test]
fn it_works() {
}
<|endoftext|>