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

//! A (mostly) lock-free concurrent work-stealing deque
//!
//! This module contains an implementation of the Chase-Lev work stealing deque
//! described in "Dynamic Circular Work-Stealing Deque". The implementation is
//! heavily based on the pseudocode found in the paper.
//!
//! This implementation does not want to have the restriction of a garbage
//! collector for reclamation of buffers, and instead it uses a shared pool of
//! buffers. This shared pool is required for correctness in this
//! implementation.
//!
//! The only lock-synchronized portions of this deque are the buffer allocation
//! and deallocation portions. Otherwise all operations are lock-free.
//!
//! # Example
//!
//!     use deque::BufferPool;
//!
//!     let mut pool = BufferPool::new();
//!     let (mut worker, mut stealer) = pool.deque();
//!
//!     // Only the worker may push/pop
//!     worker.push(1);
//!     worker.pop();
//!
//!     // Stealers take data from the other end of the deque
//!     worker.push(1);
//!     stealer.steal();
//!
//!     // Stealers can be cloned to have many stealers stealing in parallel
//!     worker.push(1);
//!     let mut stealer2 = stealer.clone();
//!     stealer2.steal();

// NB: the "buffer pool" strategy is not done for speed, but rather for
//     correctness. For more info, see the comment on `swap_buffer`

// FIXME: all atomic operations in this module use a SeqCst ordering. That is
//      probably overkill

pub use self::Stolen::*;

use std::sync::{Arc, Mutex};
use std::boxed::Box;
use std::vec::Vec;
use std::mem::{forget, size_of, transmute};
use std::ptr;
use std::slice;

use std::sync::atomic::{AtomicIsize, AtomicPtr};
use std::sync::atomic::Ordering::SeqCst;

// Once the queue is less than 1/K full, then it will be downsized. Note that
// the deque requires that this number be less than 2.
static K: isize = 4;

// Minimum number of bits that a buffer size should be. No buffer will resize to
// under this value, and all deques will initially contain a buffer of this
// size.
//
// The size in question is 1 << MIN_BITS
static MIN_BITS: usize = 7;

struct Deque<T: Send> {
    bottom: AtomicIsize,
    top: AtomicIsize,
    array: AtomicPtr<Buffer<T>>,
    pool: BufferPool<T>,
}

/// Worker half of the work-stealing deque. This worker has exclusive access to
/// one side of the deque, and uses `push` and `pop` method to manipulate it.
///
/// There may only be one worker per deque.
pub struct Worker<T: Send> {
    deque: Arc<Deque<T>>,
}

unsafe impl<T: Send> Send for Worker<T> { }

/// The stealing half of the work-stealing deque. Stealers have access to the
/// opposite end of the deque from the worker, and they only have access to the
/// `steal` method.
pub struct Stealer<T: Send> {
    deque: Arc<Deque<T>>,
}

unsafe impl<T: Send> Send for Stealer<T> { }

/// When stealing some data, this is an enumeration of the possible outcomes.
#[derive(PartialEq, Debug)]
pub enum Stolen<T> {
    /// The deque was empty at the time of stealing
    Empty,
    /// The stealer lost the race for stealing data, and a retry may return more
    /// data.
    Abort,
    /// The stealer has successfully stolen some data.
    Data(T),
}

/// The allocation pool for buffers used by work-stealing deques. Right now this
/// structure is used for reclamation of memory after it is no longer in use by
/// deques.
///
/// This data structure is protected by a mutex, but it is rarely used. Deques
/// will only use this structure when allocating a new buffer or deallocating a
/// previous one.
pub struct BufferPool<T: Send> {
    pool: Arc<Mutex<Vec<Box<Buffer<T>>>>>,
}

/// An internal buffer used by the chase-lev deque. This structure is actually
/// implemented as a circular buffer, and is used as the intermediate storage of
/// the data in the deque.
///
/// This type is implemented with *T instead of Vec<T> for two reasons:
///
///   1. There is nothing safe about using this buffer. This easily allows the
///      same value to be read twice in to rust, and there is nothing to
///      prevent this. The usage by the deque must ensure that one of the
///      values is forgotten. Furthermore, we only ever want to manually run
///      destructors for values in this buffer (on drop) because the bounds
///      are defined by the deque it's owned by.
///
///   2. We can certainly avoid bounds checks using *T instead of Vec<T>, although
///      LLVM is probably pretty good at doing this already.
struct Buffer<T: Send> {
    storage: *const T,
    log_size: usize,
}

unsafe impl<T: Send> Send for Buffer<T> { }

impl<T: Send> BufferPool<T> {
    /// Allocates a new buffer pool which in turn can be used to allocate new
    /// deques.
    pub fn new() -> BufferPool<T> {
        BufferPool { pool: Arc::new(Mutex::new(Vec::new())) }
    }

    /// Allocates a new work-stealing deque which will send/receiving memory to
    /// and from this buffer pool.
    pub fn deque(&self) -> (Worker<T>, Stealer<T>) {
        let a = Arc::new(Deque::new(self.clone()));
        let b = a.clone();
        (Worker { deque: a }, Stealer { deque: b })
    }

    fn alloc(&mut self, bits: usize) -> Box<Buffer<T>> {
        unsafe {
            let mut pool = self.pool.lock().unwrap();
            match pool.iter().position(|x| x.size() >= (1 << bits)) {
                Some(i) => pool.remove(i),
                None => Box::new(Buffer::new(bits))
            }
        }
    }

    fn free(&self, buf: Box<Buffer<T>>) {
        let mut pool = self.pool.lock().unwrap();
        match pool.iter().position(|v| v.size() > buf.size()) {
            Some(i) => pool.insert(i, buf),
            None => pool.push(buf),
        }
    }
}

impl<T: Send> Clone for BufferPool<T> {
    fn clone(&self) -> BufferPool<T> { BufferPool { pool: self.pool.clone() } }
}

impl<T: Send> Worker<T> {
    /// Pushes data onto the front of this work queue.
    pub fn push(&self, t: T) {
        unsafe { self.deque.push(t) }
    }
    /// Pops data off the front of the work queue, returning `None` on an empty
    /// queue.
    pub fn pop(&self) -> Option<T> {
        unsafe { self.deque.pop() }
    }

    /// Gets access to the buffer pool that this worker is attached to. This can
    /// be used to create more deques which share the same buffer pool as this
    /// deque.
    pub fn pool<'a>(&'a self) -> &'a BufferPool<T> {
        &self.deque.pool
    }
}

impl<T: Send> Stealer<T> {
    /// Steals work off the end of the queue (opposite of the worker's end)
    pub fn steal(&self) -> Stolen<T> {
        unsafe { self.deque.steal() }
    }

    /// Gets access to the buffer pool that this stealer is attached to. This
    /// can be used to create more deques which share the same buffer pool as
    /// this deque.
    pub fn pool<'a>(&'a self) -> &'a BufferPool<T> {
        &self.deque.pool
    }
}

impl<T: Send> Clone for Stealer<T> {
    fn clone(&self) -> Stealer<T> {
        Stealer { deque: self.deque.clone() }
    }
}

// Almost all of this code can be found directly in the paper so I'm not
// personally going to heavily comment what's going on here.

impl<T: Send> Deque<T> {
    fn new(mut pool: BufferPool<T>) -> Deque<T> {
        let buf = pool.alloc(MIN_BITS);
        Deque {
            bottom: AtomicIsize::new(0),
            top: AtomicIsize::new(0),
            array: AtomicPtr::new(unsafe { transmute(buf) }),
            pool: pool,
        }
    }

    unsafe fn push(&self, data: T) {
        let mut b = self.bottom.load(SeqCst);
        let t = self.top.load(SeqCst);
        let mut a = self.array.load(SeqCst);
        let size = b - t;
        if size >= (*a).size() - 1 {
            // You won't find this code in the chase-lev deque paper. This is
            // alluded to in a small footnote, however. We always free a buffer
            // when growing in order to prevent leaks.
            a = self.swap_buffer(b, a, (*a).resize(b, t, 1));
            b = self.bottom.load(SeqCst);
        }
        (*a).put(b, data);
        self.bottom.store(b + 1, SeqCst);
    }

    unsafe fn pop(&self) -> Option<T> {
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        let b = b - 1;
        self.bottom.store(b, SeqCst);
        let t = self.top.load(SeqCst);
        let size = b - t;
        if size < 0 {
            self.bottom.store(t, SeqCst);
            return None;
        }
        let data = (*a).get(b);
        if size > 0 {
            self.maybe_shrink(b, t);
            return Some(data);
        }
        if self.top.compare_and_swap(t, t + 1, SeqCst) == t {
            self.bottom.store(t + 1, SeqCst);
            return Some(data);
        } else {
            self.bottom.store(t + 1, SeqCst);
            forget(data); // someone else stole this value
            return None;
        }
    }

    unsafe fn steal(&self) -> Stolen<T> {
        let t = self.top.load(SeqCst);
        let old = self.array.load(SeqCst);
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        let size = b - t;
        if size <= 0 { return Empty }
        if size % (*a).size() == 0 {
            if a == old && t == self.top.load(SeqCst) {
                return Empty
            }
            return Abort
        }
        let data = (*a).get(t);
        if self.top.compare_and_swap(t, t + 1, SeqCst) == t {
            Data(data)
        } else {
            forget(data); // someone else stole this value
            Abort
        }
    }

    unsafe fn maybe_shrink(&self, b: isize, t: isize) {
        let a = self.array.load(SeqCst);
        if b - t < (*a).size() / K && b - t > (1 << MIN_BITS) {
            self.swap_buffer(b, a, (*a).resize(b, t, -1));
        }
    }

    // Helper routine not mentioned in the paper which is used in growing and
    // shrinking buffers to swap in a new buffer into place. As a bit of a
    // recap, the whole point that we need a buffer pool rather than just
    // calling malloc/free directly is that stealers can continue using buffers
    // after this method has called 'free' on it. The continued usage is simply
    // a read followed by a forget, but we must make sure that the memory can
    // continue to be read after we flag this buffer for reclamation.
    unsafe fn swap_buffer(&self, b: isize, old: *mut Buffer<T>,
                          buf: Buffer<T>) -> *mut Buffer<T> {
        let newbuf: *mut Buffer<T> = transmute(Box::new(buf));
        self.array.store(newbuf, SeqCst);
        let ss = (*newbuf).size();
        self.bottom.store(b + ss, SeqCst);
        let t = self.top.load(SeqCst);
        if self.top.compare_and_swap(t, t + ss, SeqCst) != t {
            self.bottom.store(b, SeqCst);
        }
        self.pool.free(transmute(old));
        return newbuf;
    }
}

impl<T: Send> Drop for Deque<T> {
    fn drop(&mut self) {
        let t = self.top.load(SeqCst);
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        // Free whatever is leftover in the dequeue, and then move the buffer
        // back into the pool.
        for i in t..b {
            let _: T = unsafe { (*a).get(i) };
        }
        self.pool.free(unsafe { transmute(a) });
    }
}

#[inline]
fn buffer_alloc_size<T>(log_size: usize) -> usize {
    (1 << log_size) * size_of::<T>()
}

impl<T: Send> Buffer<T> {
    unsafe fn new(log_size: usize) -> Buffer<T> {
        let size = buffer_alloc_size::<T>(log_size);

        let buffer_vec = Vec::with_capacity(size);
        let buffer: *const T = buffer_vec.get_unchecked(0);
        forget(buffer_vec);

        Buffer {
            storage: buffer as *const T,
            log_size: log_size,
        }
    }

    fn size(&self) -> isize { 1 << self.log_size }

    // Apparently LLVM cannot optimize (foo % (1 << bar)) into this implicitly
    fn mask(&self) -> isize { (1 << self.log_size) - 1 }

    unsafe fn elem(&self, i: isize) -> *const T {
        self.storage.offset(i & self.mask())
    }

    // This does not protect against loading duplicate values of the same cell,
    // nor does this clear out the contents contained within. Hence, this is a
    // very unsafe method which the caller needs to treat specially in case a
    // race is lost.
    unsafe fn get(&self, i: isize) -> T {
        ptr::read(self.elem(i))
    }

    // Unsafe because this unsafely overwrites possibly uninitialized or
    // initialized data.
    unsafe fn put(&self, i: isize, t: T) {
        ptr::write(self.elem(i) as *mut T, t);
    }

    // Again, unsafe because this has incredibly dubious ownership violations.
    // It is assumed that this buffer is immediately dropped.
    unsafe fn resize(&self, b: isize, t: isize, delta: isize) -> Buffer<T> {
        // NB: not entirely obvious, but thanks to 2's complement,
        // casting delta to usize and then adding gives the desired
        // effect.
        let buf = Buffer::new(self.log_size.wrapping_add(delta as usize));
        for i in t..b {
            buf.put(i, self.get(i));
        }
        return buf;
    }
}

impl<T: Send> Drop for Buffer<T> {
    fn drop(&mut self) {
        // It is assumed that all buffers are empty on drop.
        let size = buffer_alloc_size::<T>(self.log_size);
        unsafe { drop(slice::from_raw_parts(self.storage, size)); }
    }
}
<commit_msg>Hopefully fixed memory leak.<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A (mostly) lock-free concurrent work-stealing deque
//!
//! This module contains an implementation of the Chase-Lev work stealing deque
//! described in "Dynamic Circular Work-Stealing Deque". The implementation is
//! heavily based on the pseudocode found in the paper.
//!
//! This implementation does not want to have the restriction of a garbage
//! collector for reclamation of buffers, and instead it uses a shared pool of
//! buffers. This shared pool is required for correctness in this
//! implementation.
//!
//! The only lock-synchronized portions of this deque are the buffer allocation
//! and deallocation portions. Otherwise all operations are lock-free.
//!
//! # Example
//!
//!     use deque::BufferPool;
//!
//!     let mut pool = BufferPool::new();
//!     let (mut worker, mut stealer) = pool.deque();
//!
//!     // Only the worker may push/pop
//!     worker.push(1);
//!     worker.pop();
//!
//!     // Stealers take data from the other end of the deque
//!     worker.push(1);
//!     stealer.steal();
//!
//!     // Stealers can be cloned to have many stealers stealing in parallel
//!     worker.push(1);
//!     let mut stealer2 = stealer.clone();
//!     stealer2.steal();

// NB: the "buffer pool" strategy is not done for speed, but rather for
//     correctness. For more info, see the comment on `swap_buffer`

// FIXME: all atomic operations in this module use a SeqCst ordering. That is
//      probably overkill

pub use self::Stolen::*;

use std::sync::{Arc, Mutex};
use std::boxed::Box;
use std::vec::Vec;
use std::mem::{forget, size_of, transmute};
use std::ptr;

use std::sync::atomic::{AtomicIsize, AtomicPtr};
use std::sync::atomic::Ordering::SeqCst;

// Once the queue is less than 1/K full, then it will be downsized. Note that
// the deque requires that this number be less than 2.
static K: isize = 4;

// Minimum number of bits that a buffer size should be. No buffer will resize to
// under this value, and all deques will initially contain a buffer of this
// size.
//
// The size in question is 1 << MIN_BITS
static MIN_BITS: usize = 7;

struct Deque<T: Send> {
    bottom: AtomicIsize,
    top: AtomicIsize,
    array: AtomicPtr<Buffer<T>>,
    pool: BufferPool<T>,
}

/// Worker half of the work-stealing deque. This worker has exclusive access to
/// one side of the deque, and uses `push` and `pop` method to manipulate it.
///
/// There may only be one worker per deque.
pub struct Worker<T: Send> {
    deque: Arc<Deque<T>>,
}

unsafe impl<T: Send> Send for Worker<T> { }

/// The stealing half of the work-stealing deque. Stealers have access to the
/// opposite end of the deque from the worker, and they only have access to the
/// `steal` method.
pub struct Stealer<T: Send> {
    deque: Arc<Deque<T>>,
}

unsafe impl<T: Send> Send for Stealer<T> { }

/// When stealing some data, this is an enumeration of the possible outcomes.
#[derive(PartialEq, Debug)]
pub enum Stolen<T> {
    /// The deque was empty at the time of stealing
    Empty,
    /// The stealer lost the race for stealing data, and a retry may return more
    /// data.
    Abort,
    /// The stealer has successfully stolen some data.
    Data(T),
}

/// The allocation pool for buffers used by work-stealing deques. Right now this
/// structure is used for reclamation of memory after it is no longer in use by
/// deques.
///
/// This data structure is protected by a mutex, but it is rarely used. Deques
/// will only use this structure when allocating a new buffer or deallocating a
/// previous one.
pub struct BufferPool<T: Send> {
    pool: Arc<Mutex<Vec<Box<Buffer<T>>>>>,
}

/// An internal buffer used by the chase-lev deque. This structure is actually
/// implemented as a circular buffer, and is used as the intermediate storage of
/// the data in the deque.
///
/// This type is implemented with *T instead of Vec<T> for two reasons:
///
///   1. There is nothing safe about using this buffer. This easily allows the
///      same value to be read twice in to rust, and there is nothing to
///      prevent this. The usage by the deque must ensure that one of the
///      values is forgotten. Furthermore, we only ever want to manually run
///      destructors for values in this buffer (on drop) because the bounds
///      are defined by the deque it's owned by.
///
///   2. We can certainly avoid bounds checks using *T instead of Vec<T>, although
///      LLVM is probably pretty good at doing this already.
struct Buffer<T: Send> {
    storage: *const T,
    log_size: usize,
}

unsafe impl<T: Send> Send for Buffer<T> { }

impl<T: Send> BufferPool<T> {
    /// Allocates a new buffer pool which in turn can be used to allocate new
    /// deques.
    pub fn new() -> BufferPool<T> {
        BufferPool { pool: Arc::new(Mutex::new(Vec::new())) }
    }

    /// Allocates a new work-stealing deque which will send/receiving memory to
    /// and from this buffer pool.
    pub fn deque(&self) -> (Worker<T>, Stealer<T>) {
        let a = Arc::new(Deque::new(self.clone()));
        let b = a.clone();
        (Worker { deque: a }, Stealer { deque: b })
    }

    fn alloc(&mut self, bits: usize) -> Box<Buffer<T>> {
        unsafe {
            let mut pool = self.pool.lock().unwrap();
            match pool.iter().position(|x| x.size() >= (1 << bits)) {
                Some(i) => pool.remove(i),
                None => Box::new(Buffer::new(bits))
            }
        }
    }

    fn free(&self, buf: Box<Buffer<T>>) {
        let mut pool = self.pool.lock().unwrap();
        match pool.iter().position(|v| v.size() > buf.size()) {
            Some(i) => pool.insert(i, buf),
            None => pool.push(buf),
        }
    }
}

impl<T: Send> Clone for BufferPool<T> {
    fn clone(&self) -> BufferPool<T> { BufferPool { pool: self.pool.clone() } }
}

impl<T: Send> Worker<T> {
    /// Pushes data onto the front of this work queue.
    pub fn push(&self, t: T) {
        unsafe { self.deque.push(t) }
    }
    /// Pops data off the front of the work queue, returning `None` on an empty
    /// queue.
    pub fn pop(&self) -> Option<T> {
        unsafe { self.deque.pop() }
    }

    /// Gets access to the buffer pool that this worker is attached to. This can
    /// be used to create more deques which share the same buffer pool as this
    /// deque.
    pub fn pool<'a>(&'a self) -> &'a BufferPool<T> {
        &self.deque.pool
    }
}

impl<T: Send> Stealer<T> {
    /// Steals work off the end of the queue (opposite of the worker's end)
    pub fn steal(&self) -> Stolen<T> {
        unsafe { self.deque.steal() }
    }

    /// Gets access to the buffer pool that this stealer is attached to. This
    /// can be used to create more deques which share the same buffer pool as
    /// this deque.
    pub fn pool<'a>(&'a self) -> &'a BufferPool<T> {
        &self.deque.pool
    }
}

impl<T: Send> Clone for Stealer<T> {
    fn clone(&self) -> Stealer<T> {
        Stealer { deque: self.deque.clone() }
    }
}

// Almost all of this code can be found directly in the paper so I'm not
// personally going to heavily comment what's going on here.

impl<T: Send> Deque<T> {
    fn new(mut pool: BufferPool<T>) -> Deque<T> {
        let buf = pool.alloc(MIN_BITS);
        Deque {
            bottom: AtomicIsize::new(0),
            top: AtomicIsize::new(0),
            array: AtomicPtr::new(unsafe { transmute(buf) }),
            pool: pool,
        }
    }

    unsafe fn push(&self, data: T) {
        let mut b = self.bottom.load(SeqCst);
        let t = self.top.load(SeqCst);
        let mut a = self.array.load(SeqCst);
        let size = b - t;
        if size >= (*a).size() - 1 {
            // You won't find this code in the chase-lev deque paper. This is
            // alluded to in a small footnote, however. We always free a buffer
            // when growing in order to prevent leaks.
            a = self.swap_buffer(b, a, (*a).resize(b, t, 1));
            b = self.bottom.load(SeqCst);
        }
        (*a).put(b, data);
        self.bottom.store(b + 1, SeqCst);
    }

    unsafe fn pop(&self) -> Option<T> {
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        let b = b - 1;
        self.bottom.store(b, SeqCst);
        let t = self.top.load(SeqCst);
        let size = b - t;
        if size < 0 {
            self.bottom.store(t, SeqCst);
            return None;
        }
        let data = (*a).get(b);
        if size > 0 {
            self.maybe_shrink(b, t);
            return Some(data);
        }
        if self.top.compare_and_swap(t, t + 1, SeqCst) == t {
            self.bottom.store(t + 1, SeqCst);
            return Some(data);
        } else {
            self.bottom.store(t + 1, SeqCst);
            forget(data); // someone else stole this value
            return None;
        }
    }

    unsafe fn steal(&self) -> Stolen<T> {
        let t = self.top.load(SeqCst);
        let old = self.array.load(SeqCst);
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        let size = b - t;
        if size <= 0 { return Empty }
        if size % (*a).size() == 0 {
            if a == old && t == self.top.load(SeqCst) {
                return Empty
            }
            return Abort
        }
        let data = (*a).get(t);
        if self.top.compare_and_swap(t, t + 1, SeqCst) == t {
            Data(data)
        } else {
            forget(data); // someone else stole this value
            Abort
        }
    }

    unsafe fn maybe_shrink(&self, b: isize, t: isize) {
        let a = self.array.load(SeqCst);
        if b - t < (*a).size() / K && b - t > (1 << MIN_BITS) {
            self.swap_buffer(b, a, (*a).resize(b, t, -1));
        }
    }

    // Helper routine not mentioned in the paper which is used in growing and
    // shrinking buffers to swap in a new buffer into place. As a bit of a
    // recap, the whole point that we need a buffer pool rather than just
    // calling malloc/free directly is that stealers can continue using buffers
    // after this method has called 'free' on it. The continued usage is simply
    // a read followed by a forget, but we must make sure that the memory can
    // continue to be read after we flag this buffer for reclamation.
    unsafe fn swap_buffer(&self, b: isize, old: *mut Buffer<T>,
                          buf: Buffer<T>) -> *mut Buffer<T> {
        let newbuf: *mut Buffer<T> = transmute(Box::new(buf));
        self.array.store(newbuf, SeqCst);
        let ss = (*newbuf).size();
        self.bottom.store(b + ss, SeqCst);
        let t = self.top.load(SeqCst);
        if self.top.compare_and_swap(t, t + ss, SeqCst) != t {
            self.bottom.store(b, SeqCst);
        }
        self.pool.free(transmute(old));
        return newbuf;
    }
}

impl<T: Send> Drop for Deque<T> {
    fn drop(&mut self) {
        let t = self.top.load(SeqCst);
        let b = self.bottom.load(SeqCst);
        let a = self.array.load(SeqCst);
        // Free whatever is leftover in the dequeue, and then move the buffer
        // back into the pool.
        for i in t..b {
            let _: T = unsafe { (*a).get(i) };
        }
        self.pool.free(unsafe { transmute(a) });
    }
}

#[inline]
fn buffer_alloc_size<T>(log_size: usize) -> usize {
    (1 << log_size) * size_of::<T>()
}

impl<T: Send> Buffer<T> {
    unsafe fn new(log_size: usize) -> Buffer<T> {
        let size = buffer_alloc_size::<T>(log_size);

        let buffer_vec = Vec::with_capacity(size);
        let buffer: *const T = buffer_vec.get_unchecked(0);
        forget(buffer_vec);

        Buffer {
            storage: buffer as *const T,
            log_size: log_size,
        }
    }

    fn size(&self) -> isize { 1 << self.log_size }

    // Apparently LLVM cannot optimize (foo % (1 << bar)) into this implicitly
    fn mask(&self) -> isize { (1 << self.log_size) - 1 }

    unsafe fn elem(&self, i: isize) -> *const T {
        self.storage.offset(i & self.mask())
    }

    // This does not protect against loading duplicate values of the same cell,
    // nor does this clear out the contents contained within. Hence, this is a
    // very unsafe method which the caller needs to treat specially in case a
    // race is lost.
    unsafe fn get(&self, i: isize) -> T {
        ptr::read(self.elem(i))
    }

    // Unsafe because this unsafely overwrites possibly uninitialized or
    // initialized data.
    unsafe fn put(&self, i: isize, t: T) {
        ptr::write(self.elem(i) as *mut T, t);
    }

    // Again, unsafe because this has incredibly dubious ownership violations.
    // It is assumed that this buffer is immediately dropped.
    unsafe fn resize(&self, b: isize, t: isize, delta: isize) -> Buffer<T> {
        // NB: not entirely obvious, but thanks to 2's complement,
        // casting delta to usize and then adding gives the desired
        // effect.
        let buf = Buffer::new(self.log_size.wrapping_add(delta as usize));
        for i in t..b {
            buf.put(i, self.get(i));
        }
        return buf;
    }
}

impl<T: Send> Drop for Buffer<T> {
    fn drop(&mut self) {
        // It is assumed that all buffers are empty on drop.
        let size = buffer_alloc_size::<T>(self.log_size);
        unsafe { drop(Vec::from_raw_parts(self.storage as *mut T, 0, size)); }
    }
}
<|endoftext|>


<commit_before>use super::*;

pub trait IVim {
    fn get<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&State) -> Result<T>;
    fn update<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&mut State) -> Result<T>;
    fn loop_message<T: BufRead>(&self, input: T, languageId: Option<String>) -> Result<()>;

    /// Handle an incoming message.
    fn handle_message<H: IRpcHandler>(
        &self,
        handler: &H,
        languageId: Option<&str>,
        message: &str,
    ) -> Result<()> {
        if let Ok(output) = serde_json::from_str::<rpc::Output>(message) {
            let tx = self.update(|state| {
                state
                    .txs
                    .remove(&output.id().to_int()?)
                    .ok_or_else(|| {
                        format_err!("Failed to get channel sender! id: {:?}", output.id())
                    })?
                    .into_inner()
                    .map_err(|e| format_err!("{:?}", e))
            })?;
            let result = match output {
                rpc::Output::Success(success) => Ok(success.result),
                rpc::Output::Failure(failure) => Err(format_err!("{}", failure.error.message)),
            };
            tx.send(result)?;
            return Ok(());
        }

        // FIXME
        let message = message.replace(r#","meta":{}"#, "");

        let call = serde_json::from_str(&message)?;
        match call {
            rpc::Call::MethodCall(method_call) => {
                let result = handler.handle_request(languageId, &method_call);
                if let Err(ref err) = result {
                    if err.downcast_ref::<LCError>().is_none() {
                        error!(
                            "Error handling message. Message: {}. Error: {:?}",
                            message, result
                        );
                    }
                }
                self.output(languageId, method_call.id, result)
            }
            rpc::Call::Notification(notification) => {
                handler.handle_notification(languageId, &notification)
            }
            rpc::Call::Invalid(id) => bail!("Invalid message of id: {:?}", id),
        }
    }

    /// Send message to RPC server.
    fn write(&self, languageId: Option<&str>, message: &str) -> Result<()> {
        if let Some(languageId) = languageId {
            self.update(|state| {
                let writer = state
                    .writers
                    .get_mut(languageId)
                    .ok_or(LCError::NoLanguageServer {
                        languageId: languageId.to_owned(),
                    })?;
                write!(
                    writer,
                    "Content-Length: {}\r\n\r\n{}",
                    message.len(),
                    message
                )?;
                writer.flush()?;
                Ok(())
            })?;
        } else {
            println!("Content-Length: {}\n\n{}", message.len(), message);
        }

        Ok(())
    }

    /// Write an RPC call output.
    fn output(&self, languageId: Option<&str>, id: rpc::Id, result: Result<Value>) -> Result<()> {
        let response = match result {
            Ok(ok) => rpc::Output::Success(rpc::Success {
                jsonrpc: Some(rpc::Version::V2),
                id,
                result: ok,
            }),
            Err(err) => rpc::Output::Failure(rpc::Failure {
                jsonrpc: Some(rpc::Version::V2),
                id,
                error: err.to_rpc_error(),
            }),
        };

        let message = serde_json::to_string(&response)?;
        info!("=> {}", message);
        self.write(languageId, &message)?;
        Ok(())
    }

    /// RPC method call.
    fn call<P: Serialize, V: DeserializeOwned>(
        &self,
        languageId: Option<&str>,
        method: &str,
        params: P,
    ) -> Result<V> {
        let id = self.update(|state| {
            state.id += 1;
            Ok(state.id)
        })?;

        let method_call = rpc::MethodCall {
            jsonrpc: Some(rpc::Version::V2),
            id: rpc::Id::Num(id),
            method: method.into(),
            params: params.to_params()?,
        };

        let (tx, cx) = channel();
        self.update(|state| {
            state.txs.insert(id, Mutex::new(tx));
            Ok(())
        })?;

        let message = serde_json::to_string(&method_call)?;
        info!("=> {}", message);
        self.write(languageId, &message)?;

        let value = cx.recv_timeout(std::time::Duration::from_secs(60 * 5))??;
        Ok(serde_json::from_value(value)?)
    }

    /// RPC notification.
    fn notify<P: Serialize>(
        &self,
        languageId: Option<&str>,
        method: &str,
        params: P,
    ) -> Result<()> {
        let notification = rpc::Notification {
            jsonrpc: Some(rpc::Version::V2),
            method: method.to_owned(),
            params: params.to_params()?,
        };

        let message = serde_json::to_string(&notification)?;
        info!("=> {}", message);
        self.write(languageId, &message)?;

        Ok(())
    }

    fn eval<E, T>(&self, exp: E) -> Result<T>
    where
        E: VimExp,
        T: DeserializeOwned,
    {
        let result = self.call(None, "eval", exp.to_exp())?;
        Ok(serde_json::from_value(result)?)
    }

    fn command<S: AsRef<str>>(&self, cmd: S) -> Result<()> {
        self.call::<_, u8>(None, "execute", cmd.as_ref())?;
        Ok(())
    }

    ////// Vim builtin function wrappers ///////

    fn echo<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = escape_single_quote(message.as_ref());
        let cmd = format!("echo '{}'", message);
        self.command(cmd)
    }

    fn echo_ellipsis<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = message.as_ref().lines().collect::<Vec<_>>().join(" ");
        self.call::<_, u8>(None, "s:EchoEllipsis", message)?;
        Ok(())
    }

    fn echomsg<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echomsg '{}'", message);
        self.command(cmd)
    }

    fn echoerr<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echohl Error | echomsg '{}' | echohl None", message);
        self.command(cmd)
    }

    fn echowarn<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echohl WarningMsg | echomsg '{}' | echohl None", message);
        self.command(cmd)
    }

    fn goto_location<P: AsRef<Path>>(
        &self,
        goto_cmd: &Option<String>,
        path: P,
        line: u64,
        character: u64,
    ) -> Result<()> {
        let path = path.as_ref().to_string_lossy();
        let mut cmd = "echo | ".to_string();
        let goto;
        if let Some(ref goto_cmd) = *goto_cmd {
            goto = goto_cmd.as_str();
        } else if self.eval::<_, i64>(format!("bufnr('{}')", path))? == -1 {
            goto = "edit";
        } else {
            cmd += "execute 'normal m`' | ";
            goto = "buffer";
        };
        cmd += &format!(
            "execute 'normal! m`' | execute '{} +:call\\ cursor({},{}) ' . fnameescape('{}')",
            goto,
            line + 1,
            character + 1,
            path
        );

        self.command(cmd)
    }
}

impl IVim for Arc<RwLock<State>> {
    fn get<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&State) -> Result<T>,
    {
        let state = self.read()
            .or_else(|_| Err(err_msg("Failed to lock state")))?;
        f(&state)
    }

    fn update<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&mut State) -> Result<T>,
    {
        use log::Level;

        let mut state = RwLock::write(self).or_else(|_| Err(err_msg("Failed to lock state")))?;
        let before = if log_enabled!(Level::Debug) {
            let s = serde_json::to_string(&*state)?;
            serde_json::from_str(&s)?
        } else {
            Value::default()
        };
        let result = f(&mut state);
        let after = if log_enabled!(Level::Debug) {
            let s = serde_json::to_string(&*state)?;
            serde_json::from_str(&s)?
        } else {
            Value::default()
        };
        for (k, (v1, v2)) in diff_value(&before, &after, "state") {
            debug!("{}: {} ==> {}", k, v1, v2);
        }
        result
    }

    fn loop_message<T: BufRead>(&self, input: T, languageId: Option<String>) -> Result<()> {
        // Count how many consequent empty lines.
        let mut count_empty_lines = 0;

        let mut input = input;
        let mut content_length = 0;
        loop {
            let mut message = String::new();
            let mut line = String::new();
            if let Some(languageId) = languageId.clone() {
                input.read_line(&mut line)?;
                let line = line.trim();
                if line.is_empty() {
                    count_empty_lines += 1;
                    if count_empty_lines > 5 {
                        if let Err(err) = self.cleanup(&languageId) {
                            error!("Error when cleanup: {:?}", err);
                        }

                        let mut message =
                            format!("Language server ({}) exited unexpectedly!", languageId);
                        match get_log_server() {
                            Ok(log_server) => {
                                message += "\n\nlanguage server stderr:\n";
                                message += &log_server;
                            }
                            Err(err) => error!("Error when get_log_server: {:?}", err),
                        }
                        if let Err(err) = self.echoerr(&message) {
                            error!("Error in echoerr: {:?}", err);
                        };
                        bail!("{}", message);
                    }

                    let mut buf = vec![0; content_length];
                    input.read_exact(buf.as_mut_slice())?;
                    message = String::from_utf8(buf)?;
                } else {
                    count_empty_lines = 0;
                    if !line.starts_with("Content-Length") {
                        continue;
                    }

                    let tokens: Vec<&str> = line.splitn(2, ':').collect();
                    let len = tokens
                        .get(1)
                        .ok_or_else(|| format_err!("Failed to get length! tokens: {:?}", tokens))?
                        .trim();
                    content_length = usize::from_str(len)?;
                }
            } else if input.read_line(&mut message)? == 0 {
                break;
            }

            let message = message.trim();
            if message.is_empty() {
                continue;
            }
            info!("<= {}", message);
            let state_clone = self.clone();
            let languageId_clone = languageId.clone();
            let message_clone = message.to_string();
            let spawn_result = std::thread::Builder::new()
                .name(format!(
                    "Handler-{}",
                    languageId.as_ref().map_or("main", String::as_str)
                ))
                .spawn(move || {
                    let state = state_clone;
                    let languageId = languageId_clone;
                    let message = message_clone;

                    if let Err(err) = state.handle_message(&state, languageId.as_deref(), &message)
                    {
                        if err.downcast_ref::<LCError>().is_none() {
                            error!(
                                "Error handling message. Message: {}. Error: {:?}",
                                message, err
                            );
                        }
                    }
                });
            if let Err(err) = spawn_result {
                error!("Failed to spawn handler: {:?}", err);
            }
        }

        Ok(())
    }
}
<commit_msg>Clean error reporting.<commit_after>use super::*;

pub trait IVim {
    fn get<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&State) -> Result<T>;
    fn update<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&mut State) -> Result<T>;
    fn loop_message<T: BufRead>(&self, input: T, languageId: Option<String>) -> Result<()>;

    /// Handle an incoming message.
    fn handle_message<H: IRpcHandler>(
        &self,
        handler: &H,
        languageId: Option<&str>,
        message: &str,
    ) -> Result<()> {
        if let Ok(output) = serde_json::from_str::<rpc::Output>(message) {
            let tx = self.update(|state| {
                state
                    .txs
                    .remove(&output.id().to_int()?)
                    .ok_or_else(|| {
                        format_err!("Failed to get channel sender! id: {:?}", output.id())
                    })?
                    .into_inner()
                    .map_err(|e| format_err!("{:?}", e))
            })?;
            let result = match output {
                rpc::Output::Success(success) => Ok(success.result),
                rpc::Output::Failure(failure) => Err(format_err!("{}", failure.error.message)),
            };
            tx.send(result)?;
            return Ok(());
        }

        // FIXME
        let message = message.replace(r#","meta":{}"#, "");

        let call = serde_json::from_str(&message)?;
        match call {
            rpc::Call::MethodCall(method_call) => {
                let result = handler.handle_request(languageId, &method_call);
                if let Err(ref err) = result {
                    if err.downcast_ref::<LCError>().is_none() {
                        error!(
                            "Error handling message: {}\nMessage: {}\nError: {:?}",
                            err, message, err
                        );
                    }
                }
                self.output(languageId, method_call.id, result)
            }
            rpc::Call::Notification(notification) => {
                handler.handle_notification(languageId, &notification)
            }
            rpc::Call::Invalid(id) => bail!("Invalid message of id: {:?}", id),
        }
    }

    /// Send message to RPC server.
    fn write(&self, languageId: Option<&str>, message: &str) -> Result<()> {
        if let Some(languageId) = languageId {
            self.update(|state| {
                let writer = state
                    .writers
                    .get_mut(languageId)
                    .ok_or(LCError::NoLanguageServer {
                        languageId: languageId.to_owned(),
                    })?;
                write!(
                    writer,
                    "Content-Length: {}\r\n\r\n{}",
                    message.len(),
                    message
                )?;
                writer.flush()?;
                Ok(())
            })?;
        } else {
            println!("Content-Length: {}\n\n{}", message.len(), message);
        }

        Ok(())
    }

    /// Write an RPC call output.
    fn output(&self, languageId: Option<&str>, id: rpc::Id, result: Result<Value>) -> Result<()> {
        let response = match result {
            Ok(ok) => rpc::Output::Success(rpc::Success {
                jsonrpc: Some(rpc::Version::V2),
                id,
                result: ok,
            }),
            Err(err) => rpc::Output::Failure(rpc::Failure {
                jsonrpc: Some(rpc::Version::V2),
                id,
                error: err.to_rpc_error(),
            }),
        };

        let message = serde_json::to_string(&response)?;
        info!("=> {}", message);
        self.write(languageId, &message)?;
        Ok(())
    }

    /// RPC method call.
    fn call<P: Serialize, V: DeserializeOwned>(
        &self,
        languageId: Option<&str>,
        method: &str,
        params: P,
    ) -> Result<V> {
        let id = self.update(|state| {
            state.id += 1;
            Ok(state.id)
        })?;

        let method_call = rpc::MethodCall {
            jsonrpc: Some(rpc::Version::V2),
            id: rpc::Id::Num(id),
            method: method.into(),
            params: params.to_params()?,
        };

        let (tx, cx) = channel();
        self.update(|state| {
            state.txs.insert(id, Mutex::new(tx));
            Ok(())
        })?;

        let message = serde_json::to_string(&method_call)?;
        info!("=> {}", message);
        self.write(languageId, &message)?;

        let value = cx.recv_timeout(std::time::Duration::from_secs(60 * 5))??;
        Ok(serde_json::from_value(value)?)
    }

    /// RPC notification.
    fn notify<P: Serialize>(
        &self,
        languageId: Option<&str>,
        method: &str,
        params: P,
    ) -> Result<()> {
        let notification = rpc::Notification {
            jsonrpc: Some(rpc::Version::V2),
            method: method.to_owned(),
            params: params.to_params()?,
        };

        let message = serde_json::to_string(&notification)?;
        info!("=> {}", message);
        self.write(languageId, &message)?;

        Ok(())
    }

    fn eval<E, T>(&self, exp: E) -> Result<T>
    where
        E: VimExp,
        T: DeserializeOwned,
    {
        let result = self.call(None, "eval", exp.to_exp())?;
        Ok(serde_json::from_value(result)?)
    }

    fn command<S: AsRef<str>>(&self, cmd: S) -> Result<()> {
        self.call::<_, u8>(None, "execute", cmd.as_ref())?;
        Ok(())
    }

    ////// Vim builtin function wrappers ///////

    fn echo<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = escape_single_quote(message.as_ref());
        let cmd = format!("echo '{}'", message);
        self.command(cmd)
    }

    fn echo_ellipsis<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = message.as_ref().lines().collect::<Vec<_>>().join(" ");
        self.call::<_, u8>(None, "s:EchoEllipsis", message)?;
        Ok(())
    }

    fn echomsg<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echomsg '{}'", message);
        self.command(cmd)
    }

    fn echoerr<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echohl Error | echomsg '{}' | echohl None", message);
        self.command(cmd)
    }

    fn echowarn<S: AsRef<str>>(&self, message: S) -> Result<()> {
        let message = escape_single_quote(message);
        let cmd = format!("echohl WarningMsg | echomsg '{}' | echohl None", message);
        self.command(cmd)
    }

    fn goto_location<P: AsRef<Path>>(
        &self,
        goto_cmd: &Option<String>,
        path: P,
        line: u64,
        character: u64,
    ) -> Result<()> {
        let path = path.as_ref().to_string_lossy();
        let mut cmd = "echo | ".to_string();
        let goto;
        if let Some(ref goto_cmd) = *goto_cmd {
            goto = goto_cmd.as_str();
        } else if self.eval::<_, i64>(format!("bufnr('{}')", path))? == -1 {
            goto = "edit";
        } else {
            cmd += "execute 'normal m`' | ";
            goto = "buffer";
        };
        cmd += &format!(
            "execute 'normal! m`' | execute '{} +:call\\ cursor({},{}) ' . fnameescape('{}')",
            goto,
            line + 1,
            character + 1,
            path
        );

        self.command(cmd)
    }
}

impl IVim for Arc<RwLock<State>> {
    fn get<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&State) -> Result<T>,
    {
        let state = self.read()
            .or_else(|_| Err(err_msg("Failed to lock state")))?;
        f(&state)
    }

    fn update<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&mut State) -> Result<T>,
    {
        use log::Level;

        let mut state = RwLock::write(self).or_else(|_| Err(err_msg("Failed to lock state")))?;
        let before = if log_enabled!(Level::Debug) {
            let s = serde_json::to_string(&*state)?;
            serde_json::from_str(&s)?
        } else {
            Value::default()
        };
        let result = f(&mut state);
        let after = if log_enabled!(Level::Debug) {
            let s = serde_json::to_string(&*state)?;
            serde_json::from_str(&s)?
        } else {
            Value::default()
        };
        for (k, (v1, v2)) in diff_value(&before, &after, "state") {
            debug!("{}: {} ==> {}", k, v1, v2);
        }
        result
    }

    fn loop_message<T: BufRead>(&self, input: T, languageId: Option<String>) -> Result<()> {
        // Count how many consequent empty lines.
        let mut count_empty_lines = 0;

        let mut input = input;
        let mut content_length = 0;
        loop {
            let mut message = String::new();
            let mut line = String::new();
            if let Some(languageId) = languageId.clone() {
                input.read_line(&mut line)?;
                let line = line.trim();
                if line.is_empty() {
                    count_empty_lines += 1;
                    if count_empty_lines > 5 {
                        if let Err(err) = self.cleanup(&languageId) {
                            error!("Error when cleanup: {:?}", err);
                        }

                        let mut message =
                            format!("Language server ({}) exited unexpectedly!", languageId);
                        match get_log_server() {
                            Ok(log_server) => {
                                message += "\n\nlanguage server stderr:\n";
                                message += &log_server;
                            }
                            Err(err) => error!("Error when get_log_server: {:?}", err),
                        }
                        if let Err(err) = self.echoerr(&message) {
                            error!("Error in echoerr: {:?}", err);
                        };
                        bail!("{}", message);
                    }

                    let mut buf = vec![0; content_length];
                    input.read_exact(buf.as_mut_slice())?;
                    message = String::from_utf8(buf)?;
                } else {
                    count_empty_lines = 0;
                    if !line.starts_with("Content-Length") {
                        continue;
                    }

                    let tokens: Vec<&str> = line.splitn(2, ':').collect();
                    let len = tokens
                        .get(1)
                        .ok_or_else(|| format_err!("Failed to get length! tokens: {:?}", tokens))?
                        .trim();
                    content_length = usize::from_str(len)?;
                }
            } else if input.read_line(&mut message)? == 0 {
                break;
            }

            let message = message.trim();
            if message.is_empty() {
                continue;
            }
            info!("<= {}", message);
            let state_clone = self.clone();
            let languageId_clone = languageId.clone();
            let message_clone = message.to_string();
            let spawn_result = std::thread::Builder::new()
                .name(format!(
                    "Handler-{}",
                    languageId.as_ref().map_or("main", String::as_str)
                ))
                .spawn(move || {
                    let state = state_clone;
                    let languageId = languageId_clone;
                    let message = message_clone;

                    if let Err(err) = state.handle_message(&state, languageId.as_deref(), &message)
                    {
                        if err.downcast_ref::<LCError>().is_none() {
                            error!(
                                "Error handling message: {}\nMessage: {}\nError: {:?}",
                                err, message, err
                            );
                        }
                    }
                });
            if let Err(err) = spawn_result {
                error!("Failed to spawn handler: {:?}", err);
            }
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate num_cpus;
extern crate rustc_serialize;
extern crate threadpool;
extern crate time;
extern crate tiny_http;
extern crate url;

pub use assets::match_assets;
pub use log::LogEntry;

use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::str::FromStr;
use std::sync::{Arc, Mutex};
use std::thread;
use threadpool::ThreadPool;

pub mod input;

mod assets;
mod log;
mod router;

/// An error that one of your routes can return.
///
/// This is just a convenience enum and you don't need to use it in your project
/// if you don't want to.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RouteError {
    /// Couldn't find a way to handle this request.
    NoRouteFound,

    /// The user input is wrong.
    WrongInput,

    /// The user must be logged in.
    LoginRequired,

    /// The user entered a wrong login or password.
    WrongLoginPassword,

    /// The user is logged in but shouldn't be there.
    NotAuthorized,
}

/// Starts a server and uses the given requests handler.
///
/// The request handler takes a `&Request` and must return a `Response` to send to the user.
///
/// # Common mistakes
///
/// The handler must capture its environment by value and not by reference (`'static`). If you
/// use closure, don't forget to put `move` in front of the closure.
///
/// The handler must also be thread-safe (`Send` and `Sync`).
/// For example this handler isn't thread-safe:
///
/// ```ignore
/// let mut requests_counter = 0;
///
/// rouille::start_server("localhost:80", move |request| {
///     requests_counter += 1;
///
///     // rest of the handler
/// })
/// ```
///
/// Multiple requests can be processed simultaneously, therefore you can't mutably access
/// variables from the outside.
///
/// Instead you must use a `Mutex`:
///
/// ```no_run
/// use std::sync::Mutex;
/// let requests_counter = Mutex::new(0);
///
/// rouille::start_server("localhost:80", move |request| {
///     *requests_counter.lock().unwrap() += 1;
///
///     // rest of the handler
/// # panic!()
/// })
/// ```
///
/// # Panic handling
///
/// If your request handler panicks, a 500 error will automatically be sent to the client.
///
pub fn start_server<A, F>(addr: A, handler: F) -> !
                          where A: ToSocketAddrs,
                                F: Send + Sync + 'static + Fn(&Request) -> Response
{
    // FIXME: directly pass `addr` to the `TcpListener`
    let port = addr.to_socket_addrs().unwrap().next().unwrap().port();
    let server = tiny_http::ServerBuilder::new().with_port(port).build().unwrap();

    let handler = Arc::new(handler);
    let pool = ThreadPool::new(num_cpus::get());

    loop {
        let request = server.recv().unwrap();

        let request = Request {
            url: request.url().to_owned(),
            method: request.method().as_str().to_owned(),
            https: false,
            data: Mutex::new(None),
            inner: RequestImpl::Real(Mutex::new(Some(request))),
        };

        let handler = handler.clone();

        pool.execute(move || {
            let response = (*handler)(&request);
            match request.inner {
                RequestImpl::Real(rq) => rq.lock().unwrap().take().unwrap()
                                           .respond(response.response),
                RequestImpl::Fake { .. } => unreachable!()
            }
        });
    }
}

/// Represents a request made by the client.
pub struct Request {
    url: String,
    method: String,
    https: bool,
    data: Mutex<Option<Vec<u8>>>,
    inner: RequestImpl,
}

enum RequestImpl {
    Real(Mutex<Option<tiny_http::Request>>),     // TODO: when Mutex gets "into_inner", remove the Option
    Fake { headers: Vec<(String, String)> },
}

impl Request {
    /// Builds a fake request to be used during tests.
    pub fn fake<U, M>(https: bool, url: U, method: M, headers: Vec<(String, String)>, data: Vec<u8>)
                      -> Request where U: Into<String>, M: Into<String>
    {
        Request {
            url: url.into(),
            method: method.into(),
            https: https,
            data: Mutex::new(Some(data)),
            inner: RequestImpl::Fake { headers: headers },
        }
    }

    /// Returns `true` if the request uses HTTPS instead of HTTP.
    #[inline]
    pub fn secure(&self) -> bool {
        self.https
    }

    /// Returns the method of the request (`GET`, `POST`, etc.).
    #[inline]
    pub fn method(&self) -> &str {
        &self.method
    }

    /// Returns the URL requested by the client. It is not decoded and thus can contain `%20` or
    /// other characters.
    #[inline]
    pub fn url(&self) -> &str {
        &self.url
    }

    /// Returns the value of a header of the request.
    #[inline]
    pub fn header(&self, key: &str) -> Option<String> {
        match self.inner {
            RequestImpl::Real(ref request) => request.lock().unwrap().as_ref().unwrap().headers()
                                                     .iter().find(|h| h.field.as_str() == key)
                                                     .map(|h| h.value.as_str().to_owned()),

            RequestImpl::Fake { ref headers } => headers.iter().find(|&&(ref k, _)| k == key)
                                                        .map(|&(_, ref v)| v.clone())
        }
    }

    /// Returns the data of the request.
    pub fn data(&self) -> Vec<u8> {
        let mut data = self.data.lock().unwrap();

        if let Some(ref mut data) = *data {
            return data.clone();
        }

        let mut read = Vec::new();
        match self.inner {
            RequestImpl::Real(ref request) => {
                let _ = request.lock().unwrap().as_mut().unwrap()
                               .as_reader().read_to_end(&mut read);
            },
            RequestImpl::Fake { .. } => ()
        };
        let read_clone = read.clone();
        *data = Some(read);
        read_clone
    }
}

/// Contains a prototype of a response. The response is only sent when you call `Request::respond`.
pub struct Response {
    response: tiny_http::ResponseBox,
}

impl Response {
    /// UNSTABLE. Builds a default response to handle the given route error.
    ///
    /// Important: don't use this in a real website. This function is just a convenience when
    /// prototyping.
    ///
    /// For authentication-related errors, you are strongly encouraged to handle them yourself.
    #[inline]
    pub fn from_error(err: &RouteError) -> Response {
        let response = match err {
            &RouteError::NoRouteFound => tiny_http::Response::empty(404),
            &RouteError::WrongInput => tiny_http::Response::empty(400),
            &RouteError::LoginRequired => tiny_http::Response::empty(401),     // TODO: www-auth header?
            &RouteError::WrongLoginPassword => tiny_http::Response::empty(401),     // TODO: www-auth header?
            &RouteError::NotAuthorized => tiny_http::Response::empty(403),
        };

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that redirects the user to another URL.
    #[inline]
    pub fn redirect(target: &str) -> Response {
        let response = tiny_http::Response::empty(303);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str(&format!("Location: {}", target)).unwrap());

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that outputs HTML.
    #[inline]
    pub fn html(content: &[u8]) -> Response {
        let response = tiny_http::Response::from_data(content);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str("Content-Type: text/html; charset=utf8").unwrap());

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that outputs JSON.
    #[inline]
    pub fn json<T>(content: &T) -> Response where T: rustc_serialize::Encodable {
        let data = rustc_serialize::json::encode(content).unwrap();
        let response = tiny_http::Response::from_string(data);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str("Content-Type: application/json").unwrap());

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that returns a `401 Not Authorized` status
    /// and a `WWW-Authenticate` header.
    #[inline]
    pub fn basic_http_auth_login_required(realm: &str) -> Response {
        // TODO: escape the realm
        let response = tiny_http::Response::empty(401);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str(&format!("WWW-Authenticate: Basic realm=\"{}\"", realm)).unwrap());

        Response { response: response.boxed() }
    }
}
<commit_msg>Add Response::text<commit_after>extern crate num_cpus;
extern crate rustc_serialize;
extern crate threadpool;
extern crate time;
extern crate tiny_http;
extern crate url;

pub use assets::match_assets;
pub use log::LogEntry;

use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::str::FromStr;
use std::sync::{Arc, Mutex};
use std::thread;
use threadpool::ThreadPool;

pub mod input;

mod assets;
mod log;
mod router;

/// An error that one of your routes can return.
///
/// This is just a convenience enum and you don't need to use it in your project
/// if you don't want to.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RouteError {
    /// Couldn't find a way to handle this request.
    NoRouteFound,

    /// The user input is wrong.
    WrongInput,

    /// The user must be logged in.
    LoginRequired,

    /// The user entered a wrong login or password.
    WrongLoginPassword,

    /// The user is logged in but shouldn't be there.
    NotAuthorized,
}

/// Starts a server and uses the given requests handler.
///
/// The request handler takes a `&Request` and must return a `Response` to send to the user.
///
/// # Common mistakes
///
/// The handler must capture its environment by value and not by reference (`'static`). If you
/// use closure, don't forget to put `move` in front of the closure.
///
/// The handler must also be thread-safe (`Send` and `Sync`).
/// For example this handler isn't thread-safe:
///
/// ```ignore
/// let mut requests_counter = 0;
///
/// rouille::start_server("localhost:80", move |request| {
///     requests_counter += 1;
///
///     // rest of the handler
/// })
/// ```
///
/// Multiple requests can be processed simultaneously, therefore you can't mutably access
/// variables from the outside.
///
/// Instead you must use a `Mutex`:
///
/// ```no_run
/// use std::sync::Mutex;
/// let requests_counter = Mutex::new(0);
///
/// rouille::start_server("localhost:80", move |request| {
///     *requests_counter.lock().unwrap() += 1;
///
///     // rest of the handler
/// # panic!()
/// })
/// ```
///
/// # Panic handling
///
/// If your request handler panicks, a 500 error will automatically be sent to the client.
///
pub fn start_server<A, F>(addr: A, handler: F) -> !
                          where A: ToSocketAddrs,
                                F: Send + Sync + 'static + Fn(&Request) -> Response
{
    // FIXME: directly pass `addr` to the `TcpListener`
    let port = addr.to_socket_addrs().unwrap().next().unwrap().port();
    let server = tiny_http::ServerBuilder::new().with_port(port).build().unwrap();

    let handler = Arc::new(handler);
    let pool = ThreadPool::new(num_cpus::get());

    loop {
        let request = server.recv().unwrap();

        let request = Request {
            url: request.url().to_owned(),
            method: request.method().as_str().to_owned(),
            https: false,
            data: Mutex::new(None),
            inner: RequestImpl::Real(Mutex::new(Some(request))),
        };

        let handler = handler.clone();

        pool.execute(move || {
            let response = (*handler)(&request);
            match request.inner {
                RequestImpl::Real(rq) => rq.lock().unwrap().take().unwrap()
                                           .respond(response.response),
                RequestImpl::Fake { .. } => unreachable!()
            }
        });
    }
}

/// Represents a request made by the client.
pub struct Request {
    url: String,
    method: String,
    https: bool,
    data: Mutex<Option<Vec<u8>>>,
    inner: RequestImpl,
}

enum RequestImpl {
    Real(Mutex<Option<tiny_http::Request>>),     // TODO: when Mutex gets "into_inner", remove the Option
    Fake { headers: Vec<(String, String)> },
}

impl Request {
    /// Builds a fake request to be used during tests.
    pub fn fake<U, M>(https: bool, url: U, method: M, headers: Vec<(String, String)>, data: Vec<u8>)
                      -> Request where U: Into<String>, M: Into<String>
    {
        Request {
            url: url.into(),
            method: method.into(),
            https: https,
            data: Mutex::new(Some(data)),
            inner: RequestImpl::Fake { headers: headers },
        }
    }

    /// Returns `true` if the request uses HTTPS instead of HTTP.
    #[inline]
    pub fn secure(&self) -> bool {
        self.https
    }

    /// Returns the method of the request (`GET`, `POST`, etc.).
    #[inline]
    pub fn method(&self) -> &str {
        &self.method
    }

    /// Returns the URL requested by the client. It is not decoded and thus can contain `%20` or
    /// other characters.
    #[inline]
    pub fn url(&self) -> &str {
        &self.url
    }

    /// Returns the value of a header of the request.
    #[inline]
    pub fn header(&self, key: &str) -> Option<String> {
        match self.inner {
            RequestImpl::Real(ref request) => request.lock().unwrap().as_ref().unwrap().headers()
                                                     .iter().find(|h| h.field.as_str() == key)
                                                     .map(|h| h.value.as_str().to_owned()),

            RequestImpl::Fake { ref headers } => headers.iter().find(|&&(ref k, _)| k == key)
                                                        .map(|&(_, ref v)| v.clone())
        }
    }

    /// Returns the data of the request.
    pub fn data(&self) -> Vec<u8> {
        let mut data = self.data.lock().unwrap();

        if let Some(ref mut data) = *data {
            return data.clone();
        }

        let mut read = Vec::new();
        match self.inner {
            RequestImpl::Real(ref request) => {
                let _ = request.lock().unwrap().as_mut().unwrap()
                               .as_reader().read_to_end(&mut read);
            },
            RequestImpl::Fake { .. } => ()
        };
        let read_clone = read.clone();
        *data = Some(read);
        read_clone
    }
}

/// Contains a prototype of a response. The response is only sent when you call `Request::respond`.
pub struct Response {
    response: tiny_http::ResponseBox,
}

impl Response {
    /// UNSTABLE. Builds a default response to handle the given route error.
    ///
    /// Important: don't use this in a real website. This function is just a convenience when
    /// prototyping.
    ///
    /// For authentication-related errors, you are strongly encouraged to handle them yourself.
    #[inline]
    pub fn from_error(err: &RouteError) -> Response {
        let response = match err {
            &RouteError::NoRouteFound => tiny_http::Response::empty(404),
            &RouteError::WrongInput => tiny_http::Response::empty(400),
            &RouteError::LoginRequired => tiny_http::Response::empty(401),     // TODO: www-auth header?
            &RouteError::WrongLoginPassword => tiny_http::Response::empty(401),     // TODO: www-auth header?
            &RouteError::NotAuthorized => tiny_http::Response::empty(403),
        };

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that redirects the user to another URL.
    #[inline]
    pub fn redirect(target: &str) -> Response {
        let response = tiny_http::Response::empty(303);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str(&format!("Location: {}", target)).unwrap());

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that outputs HTML.
    #[inline]
    pub fn html(content: &[u8]) -> Response {
        let response = tiny_http::Response::from_data(content);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str("Content-Type: text/html; charset=utf8").unwrap());

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that outputs plain text.
    #[inline]
    pub fn text(text: &str) -> Response {
        let response = tiny_http::Response::from_string(text);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str("Content-Type: text/plain; charset=utf8").unwrap());

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that outputs JSON.
    #[inline]
    pub fn json<T>(content: &T) -> Response where T: rustc_serialize::Encodable {
        let data = rustc_serialize::json::encode(content).unwrap();
        let response = tiny_http::Response::from_string(data);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str("Content-Type: application/json").unwrap());

        Response { response: response.boxed() }
    }

    /// Builds a `Response` that returns a `401 Not Authorized` status
    /// and a `WWW-Authenticate` header.
    #[inline]
    pub fn basic_http_auth_login_required(realm: &str) -> Response {
        // TODO: escape the realm
        let response = tiny_http::Response::empty(401);
        // TODO: slow \|/
        let response = response.with_header(tiny_http::Header::from_str(&format!("WWW-Authenticate: Basic realm=\"{}\"", realm)).unwrap());

        Response { response: response.boxed() }
    }
}
<|endoftext|>


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

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

mod common;
pub use common::*;

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

    static __afl_fuzz_len: *const u32;
    static __afl_fuzz_ptr: *const u8;
}

#[no_mangle]
pub static __afl_sharedmem_fuzzing: i32 = 1;

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

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

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

    let mut input = vec![];

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

    while unsafe { __afl_persistent_loop(1000) } != 0 {
        // get the testcase from the fuzzer
        let input_ref = if unsafe { __afl_fuzz_ptr.is_null() } {
            // in-memory testcase delivery is not enabled
            // get buffer from AFL++ through stdin
            let result = io::stdin().read_to_end(&mut input);
            if result.is_err() {
                return;
            }
            &input
        } else {
            unsafe {
                // get the testcase from the shared memory
                let input_len = *__afl_fuzz_len as usize;
                std::slice::from_raw_parts(__afl_fuzz_ptr, input_len)
            }
        };

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

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

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

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

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

            $body
        });
    };
}

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

    extern crate libc;
    extern crate tempdir;

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

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

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

mod common;
pub use common::*;

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

    static __afl_fuzz_len: *const u32;
    static __afl_fuzz_ptr: *const u8;
}

#[no_mangle]
pub static __afl_sharedmem_fuzzing: i32 = 1;

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

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

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

    let mut input = vec![];

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

    while unsafe { __afl_persistent_loop(1000) } != 0 {
        // get the testcase from the fuzzer
        let input_ref = if unsafe { __afl_fuzz_ptr.is_null() } {
            // in-memory testcase delivery is not enabled
            // get buffer from AFL++ through stdin
            let result = io::stdin().read_to_end(&mut input);
            if result.is_err() {
                return;
            }
            &input
        } else {
            unsafe {
                // get the testcase from the shared memory
                let input_len = *__afl_fuzz_len as usize;
                std::slice::from_raw_parts(__afl_fuzz_ptr, input_len)
            }
        };

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

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

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

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

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

            $body
        });
    };
}

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

    extern crate libc;
    extern crate tempdir;

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

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


<commit_before>//! A [TOML]-parsing library
//!
//! This library implements a [TOML] v0.5.0 compatible parser,
//! primarily supporting the [`serde`] library for encoding/decoding
//! various types in Rust.
//!
//! TOML itself is a simple, ergonomic, and readable configuration format:
//!
//! ```toml
//! [package]
//! name = "toml"
//! version = "0.4.2"
//! authors = ["Alex Crichton <alex@alexcrichton.com>"]
//!
//! [dependencies]
//! serde = "1.0"
//! ```
//!
//! The TOML format tends to be relatively common throughout the Rust community
//! for configuration, notably being used by [Cargo], Rust's package manager.
//!
//! ## TOML values
//!
//! A value in TOML is represented with the `Value` enum in this crate:
//!
//! ```rust,ignore
//! pub enum Value {
//!     String(String),
//!     Integer(i64),
//!     Float(f64),
//!     Boolean(bool),
//!     Datetime(Datetime),
//!     Array(Array),
//!     Table(Table),
//! }
//! ```
//!
//! TOML is similar to JSON with the notable addition of a `Datetime`
//! type. In general, TOML and JSON are interchangeable in terms of
//! formats.
//!
//! ## Parsing TOML
//!
//! The easiest way to parse a TOML document is via the `Value` type:
//!
//! ```rust
//! use toml::Value;
//!
//! let value = "foo = 'bar'".parse::<Value>().unwrap();
//!
//! assert_eq!(value["foo"].as_str(), Some("bar"));
//! ```
//!
//! The `Value` type implements a number of convenience methods and
//! traits; the example above uses `FromStr` to parse a `str` into a
//! `Value`.
//!
//! ## Deserialization and Serialization
//!
//! This crate supports [`serde`] 1.0 with a number of
//! implementations of the `Deserialize`, `Serialize`, `Deserializer`, and
//! `Serializer` traits. Namely, you'll find:
//!
//! * `Deserialize for Value`
//! * `Serialize for Value`
//! * `Deserialize for Datetime`
//! * `Serialize for Datetime`
//! * `Deserializer for de::Deserializer`
//! * `Serializer for ser::Serializer`
//! * `Deserializer for Value`
//!
//! This means that you can use Serde to deserialize/serialize the
//! `Value` type as well as the `Datetime` type in this crate. You can also
//! use the `Deserializer`, `Serializer`, or `Value` type itself to act as
//! a deserializer/serializer for arbitrary types.
//!
//! An example of deserializing with TOML is:
//!
//! ```rust
//! use serde_derive::Deserialize;
//!
//! #[derive(Deserialize)]
//! struct Config {
//!     ip: String,
//!     port: Option<u16>,
//!     keys: Keys,
//! }
//!
//! #[derive(Deserialize)]
//! struct Keys {
//!     github: String,
//!     travis: Option<String>,
//! }
//!
//! fn main() {
//!     let config: Config = toml::from_str(r#"
//!         ip = '127.0.0.1'
//!
//!         [keys]
//!         github = 'xxxxxxxxxxxxxxxxx'
//!         travis = 'yyyyyyyyyyyyyyyyy'
//!     "#).unwrap();
//!
//!     assert_eq!(config.ip, "127.0.0.1");
//!     assert_eq!(config.port, None);
//!     assert_eq!(config.keys.github, "xxxxxxxxxxxxxxxxx");
//!     assert_eq!(config.keys.travis.as_ref().unwrap(), "yyyyyyyyyyyyyyyyy");
//! }
//! ```
//!
//! You can serialize types in a similar fashion:
//!
//! ```rust
//! use serde_derive::Serialize;
//!
//! #[derive(Serialize)]
//! struct Config {
//!     ip: String,
//!     port: Option<u16>,
//!     keys: Keys,
//! }
//!
//! #[derive(Serialize)]
//! struct Keys {
//!     github: String,
//!     travis: Option<String>,
//! }
//!
//! fn main() {
//!     let config = Config {
//!         ip: "127.0.0.1".to_string(),
//!         port: None,
//!         keys: Keys {
//!             github: "xxxxxxxxxxxxxxxxx".to_string(),
//!             travis: Some("yyyyyyyyyyyyyyyyy".to_string()),
//!         },
//!     };
//!
//!     let toml = toml::to_string(&config).unwrap();
//! }
//! ```
//!
//! [TOML]: https://github.com/toml-lang/toml
//! [Cargo]: https://crates.io/
//! [`serde`]: https://serde.rs/

#![doc(html_root_url = "https://docs.rs/toml/0.5")]
#![deny(missing_docs)]
#![warn(rust_2018_idioms)]

pub mod map;
pub mod value;
#[doc(no_inline)]
pub use crate::value::Value;
mod datetime;

pub mod ser;
#[doc(no_inline)]
pub use crate::ser::{to_string, to_string_pretty, to_vec, Serializer};
pub mod de;
#[doc(no_inline)]
pub use crate::de::{from_slice, from_str, Deserializer};
mod tokens;

#[doc(hidden)]
pub mod macros;

mod spanned;
pub use crate::spanned::Spanned;
<commit_msg>Forbid unsafe code (#332)<commit_after>//! A [TOML]-parsing library
//!
//! This library implements a [TOML] v0.5.0 compatible parser,
//! primarily supporting the [`serde`] library for encoding/decoding
//! various types in Rust.
//!
//! TOML itself is a simple, ergonomic, and readable configuration format:
//!
//! ```toml
//! [package]
//! name = "toml"
//! version = "0.4.2"
//! authors = ["Alex Crichton <alex@alexcrichton.com>"]
//!
//! [dependencies]
//! serde = "1.0"
//! ```
//!
//! The TOML format tends to be relatively common throughout the Rust community
//! for configuration, notably being used by [Cargo], Rust's package manager.
//!
//! ## TOML values
//!
//! A value in TOML is represented with the `Value` enum in this crate:
//!
//! ```rust,ignore
//! pub enum Value {
//!     String(String),
//!     Integer(i64),
//!     Float(f64),
//!     Boolean(bool),
//!     Datetime(Datetime),
//!     Array(Array),
//!     Table(Table),
//! }
//! ```
//!
//! TOML is similar to JSON with the notable addition of a `Datetime`
//! type. In general, TOML and JSON are interchangeable in terms of
//! formats.
//!
//! ## Parsing TOML
//!
//! The easiest way to parse a TOML document is via the `Value` type:
//!
//! ```rust
//! use toml::Value;
//!
//! let value = "foo = 'bar'".parse::<Value>().unwrap();
//!
//! assert_eq!(value["foo"].as_str(), Some("bar"));
//! ```
//!
//! The `Value` type implements a number of convenience methods and
//! traits; the example above uses `FromStr` to parse a `str` into a
//! `Value`.
//!
//! ## Deserialization and Serialization
//!
//! This crate supports [`serde`] 1.0 with a number of
//! implementations of the `Deserialize`, `Serialize`, `Deserializer`, and
//! `Serializer` traits. Namely, you'll find:
//!
//! * `Deserialize for Value`
//! * `Serialize for Value`
//! * `Deserialize for Datetime`
//! * `Serialize for Datetime`
//! * `Deserializer for de::Deserializer`
//! * `Serializer for ser::Serializer`
//! * `Deserializer for Value`
//!
//! This means that you can use Serde to deserialize/serialize the
//! `Value` type as well as the `Datetime` type in this crate. You can also
//! use the `Deserializer`, `Serializer`, or `Value` type itself to act as
//! a deserializer/serializer for arbitrary types.
//!
//! An example of deserializing with TOML is:
//!
//! ```rust
//! use serde_derive::Deserialize;
//!
//! #[derive(Deserialize)]
//! struct Config {
//!     ip: String,
//!     port: Option<u16>,
//!     keys: Keys,
//! }
//!
//! #[derive(Deserialize)]
//! struct Keys {
//!     github: String,
//!     travis: Option<String>,
//! }
//!
//! fn main() {
//!     let config: Config = toml::from_str(r#"
//!         ip = '127.0.0.1'
//!
//!         [keys]
//!         github = 'xxxxxxxxxxxxxxxxx'
//!         travis = 'yyyyyyyyyyyyyyyyy'
//!     "#).unwrap();
//!
//!     assert_eq!(config.ip, "127.0.0.1");
//!     assert_eq!(config.port, None);
//!     assert_eq!(config.keys.github, "xxxxxxxxxxxxxxxxx");
//!     assert_eq!(config.keys.travis.as_ref().unwrap(), "yyyyyyyyyyyyyyyyy");
//! }
//! ```
//!
//! You can serialize types in a similar fashion:
//!
//! ```rust
//! use serde_derive::Serialize;
//!
//! #[derive(Serialize)]
//! struct Config {
//!     ip: String,
//!     port: Option<u16>,
//!     keys: Keys,
//! }
//!
//! #[derive(Serialize)]
//! struct Keys {
//!     github: String,
//!     travis: Option<String>,
//! }
//!
//! fn main() {
//!     let config = Config {
//!         ip: "127.0.0.1".to_string(),
//!         port: None,
//!         keys: Keys {
//!             github: "xxxxxxxxxxxxxxxxx".to_string(),
//!             travis: Some("yyyyyyyyyyyyyyyyy".to_string()),
//!         },
//!     };
//!
//!     let toml = toml::to_string(&config).unwrap();
//! }
//! ```
//!
//! [TOML]: https://github.com/toml-lang/toml
//! [Cargo]: https://crates.io/
//! [`serde`]: https://serde.rs/

#![doc(html_root_url = "https://docs.rs/toml/0.5")]
#![deny(missing_docs)]
#![warn(rust_2018_idioms)]
// Makes rustc abort compilation if there are any unsafe blocks in the crate.
// Presence of this annotation is picked up by tools such as cargo-geiger
// and lets them ensure that there is indeed no unsafe code as opposed to
// something they couldn't detect (e.g. unsafe added via macro expansion, etc).
#![forbid(unsafe_code)]

pub mod map;
pub mod value;
#[doc(no_inline)]
pub use crate::value::Value;
mod datetime;

pub mod ser;
#[doc(no_inline)]
pub use crate::ser::{to_string, to_string_pretty, to_vec, Serializer};
pub mod de;
#[doc(no_inline)]
pub use crate::de::{from_slice, from_str, Deserializer};
mod tokens;

#[doc(hidden)]
pub mod macros;

mod spanned;
pub use crate::spanned::Spanned;
<|endoftext|>