<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Simple buddy-block allocator

mod math;
#[cfg(feature = "buddy_as_system")]
pub mod system;

use super::Framesque;
use self::math::PowersOf2;

use core::mem;
use core::cmp::{max, min};
use core::ptr::Unique;

use intrusive::list::{List, Node};
use intrusive::rawlink::RawLink;

use memory::alloc::Allocator;
use memory::arch::PAGE_SIZE;

#[cfg(target_os = "linux")]
#[cfg(test)]
mod test;

/// A `FreeList` is a list of unique free blocks
pub type FreeList = List<Unique<FreeBlock>, FreeBlock>;

/// A free block header stores a pointer to the next and previous free blocks.
pub struct FreeBlock { next: RawLink<FreeBlock>
                     , prev: RawLink<FreeBlock>
                     }

impl Framesque for FreeBlock {
    #[inline] fn as_ptr(&self) -> *mut u8 {
        unsafe { mem::transmute(self) } // HOPEFULLY this is good
    }
}

impl Node for FreeBlock {
    #[inline] fn prev(&self) -> &RawLink<FreeBlock> {
        &self.prev
    }
    #[inline] fn next(&self) -> &RawLink<FreeBlock> {
        &self.next
    }
    #[inline] fn prev_mut(&mut self) -> &mut RawLink<FreeBlock> {
        &mut self.prev
    }
    #[inline] fn next_mut(&mut self) -> &mut RawLink<FreeBlock> {
        &mut self.next
    }
}

// Variadic macro for taking the maximum of n > 2 numbers.
// because I'm lazy.
macro_rules! max {
    ($x:expr) => ($x);
    ($x:expr, $($xs:expr),+) => (max($x, max!($($xs),+)));
}

/// Structure with data for implementing the buddy block allocation strategy.
pub struct BuddyHeapAllocator<'a> {
    /// Address of the base of the heap. This must be aligned
    /// on a `MIN_ALIGN` boundary.
    pub start_addr: *mut u8
  , /// The allocator's free list
    free_lists: &'a mut [FreeList]
  , /// Number of blocks in the heap (must be a power of 2)
    pub heap_size: usize
  , /// Minimum block size
    pub min_block_size: usize
}

impl<'a> BuddyHeapAllocator<'a> {
    /// Construct a new `BuddyHeapAllocator`.
    ///
    /// # Arguments:
    /// + `start_addr`: a pointer to the start location of the heap
    /// + `free_lists`: an array of `FreeList`s. The cardinality
    ///     of the `free_lists` array should be equal to the maximum
    ///     allocateable order.
    /// + `heap_size`: the size of the heap (in bytes)
    ///
    /// # Returns:
    /// + A new `BuddyHeapAllocator`, obviously.
    pub unsafe fn new( start_addr: *mut u8
                     , free_lists: &'a mut [FreeList]
                     , heap_size: usize)
                     -> BuddyHeapAllocator<'a> {
        // Cache the number of free lists hopefully saving performance.
        let n_free_lists = free_lists.len();

        assert!( !start_addr.is_null()
                , "Heap start address cannot be null." );
        assert!( n_free_lists > 0
               , "Allocator must have at least one free list.");
        assert!( start_addr as usize & (PAGE_SIZE-1) as usize == 0
               , "Heap start address must be aligned on a 4k boundary.");

        let min_block_size = heap_size >> (n_free_lists - 1);

        assert!( heap_size >= min_block_size
               , "Heap must be large enough to contain at least one block.");
        assert!( min_block_size >= mem::size_of::<FreeBlock>()
               , "Minimum block size must be large enough to contain \
                  the free block header.");
        assert!( heap_size.is_pow2()
               , "Heap size must be a power of 2.");

    //    // We must have one free list per possible heap block size.
    //    assert_eq!(min_block_size *
    //               (2u32.pow(free_lists.len() as u32 - 1)) as usize,
    //               heap_size);

        // Zero out the free lists in case we were passed existing data.
        for list in free_lists.iter_mut() {
            *list = FreeList::new();
        }

        let mut heap
            = BuddyHeapAllocator { start_addr: start_addr
                                 , free_lists: free_lists
                                 , heap_size: heap_size
                                 , min_block_size: min_block_size
                                 };

        // the order needed to allocate the entire heap as a single block
        let root_order
            = heap.alloc_order(heap_size, 1)
                  .expect("Couldn't determine heap root allocation order!\
                           This should be (as far as I know) impossible.\
                           Something is seriously amiss.");

        // Push the entire heap onto the free lists as the first block.
        heap.push_block(start_addr, root_order);
        heap
    }

    /// Add a block of max order
    pub unsafe fn add_block(&mut self, block: *mut u8) {
        let order = self.free_lists.len() -1;
        self.push_block(block, order);
    }

    /// Computes the size of an allocation request.
    ///
    /// # Arguments
    /// + `size`: A `usize` containing the size of the request
    /// + `align`: A `usize` containing the alignment of the request
    ///
    /// # Returns
    /// + `None` if the request is invalid
    /// + `Some(usize)` containing the size needed if the request is valid
    #[inline]
    pub fn alloc_size(&self, size: usize, align: usize) -> Option<usize> {
        // Pre-check if this is a valid allocation request:
        //  - allocations must be aligned on power of 2 boundaries
        //  - we cannot allocate requests with alignments greater than the
        //    base alignment of the heap without jumping through a bunch of
        //    hoops.
        if !align.is_pow2() || align > PAGE_SIZE as usize {
            None
        // If the request is valid, compute the size we need to allocate
        } else {
            // the allocation size for the request is the next power of 2
            // after the size of the request, the alignment of the request,
            // or the minimum block size (whichever is greatest).
            let alloc_size = max!(size, self.min_block_size, align).next_pow2();

            if alloc_size > self.heap_size {
                // if the calculated size is greater than the size of the heap,
                // we (obviously) cannot allocate this request.
                None
            } else {
                // otherwise, return the calculated size.
                Some(alloc_size)
            }
        }
    }

    /// Computes the order of an allocation request.
    ///
    /// The "order" of an allocation refers to the number of times we need to
    /// double the minimum block size to get a large enough block for that
    /// allocation.
    #[inline]
    pub fn alloc_order(&self, size: usize, align: usize) -> Option<usize> {
        trace!(target: "alloc", "TRACE: alloc_order() called");
        self.alloc_size(size, align)
            .map(|s| {
                trace!(target: "alloc"
                      , "in alloc_order(): alloc_size() returned {}", s);
                s.log2() - self.min_block_size.log2()
            })
    }

    /// Computes the size  allocated for a request of the given order.
    #[inline]
    fn order_alloc_size(&self, order: usize) -> usize {
        1 << (self.min_block_size.log2() + order)
    }

    #[inline]
    unsafe fn push_block(&mut self, ptr: *mut u8, order: usize) {
        self.free_lists[order]
            .push_front(Unique::new(ptr as *mut FreeBlock))
    }

    #[inline]
    unsafe fn pop_block(&mut self, order: usize) -> Option<*mut u8>{
        self.free_lists[order]
            .pop_front()
            .map(|block| block.get().as_ptr())
    }

    /// Splits a block
    unsafe fn split_block( &mut self
                         , block: *mut u8
                         , old_order: usize
                         , new_order: usize ) {
        trace!( target: "alloc"
              , "split_block() was called, target order: {}.", new_order);

        assert!( new_order < old_order
               , "Cannot split a block larger than its current order!");
        assert!( old_order <= self.free_lists.len()
               , "Order is too large for this allocator!");

        let mut split_size = self.order_alloc_size(old_order);
        // let mut curr_order = order;
        for order in (new_order..old_order).rev() {
            split_size >>= 1;

            // let split_size = self.order_alloc_size(order);
            self.push_block(block.offset(split_size as isize), order);

            trace!( target: "alloc"
                  , "split block successfully, order: {}, split size: {}"
                  , order, split_size );

        }
    }

    /// Finds the buddy block for a given block.
    ///
    /// # Arguments
    /// + `order`: the order of the block to find a buddy for
    /// + `block`: a pointer to the block to find a buddy for
    ///
    /// # Returns
    /// + `Some(*mut u8)` pointing to the buddy block if a buddy was found
    /// + `None` if the block was the size of the entire heap
    pub unsafe fn get_buddy( &self
                           , order: usize
                           , block: *mut u8)
                            -> Option<*mut u8> {
        // Determine the size of the block allocated for the given order
        let block_size = self.order_alloc_size(order);
        if block_size < self.heap_size {
            // Determine the block's position in the heap.
            let block_pos = (block as usize) - (self.start_addr as usize);
            // Calculate the block's buddy by XORing the block's position
            // in the heap with its size.
            Some(self.start_addr.offset((block_pos ^ block_size) as isize))
        } else {
            // If the block is the size of the entire heap, it (obviously)
            // cannot have a buddy block.
            None
        }
    }

    /// Finds and removes the target block from the free list.
    ///
    /// # Arguments
    /// + `order`: the order of the free list to remove the block from_raw
    /// + `block`: a pointer to the block to remove
    ///
    /// # Returns
    /// + `true` if the block was found and removed from the free List
    /// + `false` if the block was not found
    pub fn remove_block(&mut self, order: usize, block: *mut u8) -> bool {
        self.free_lists[order]
            .cursor_mut()
            .find_and_remove(|b| b as *const FreeBlock as *const u8 == block)
            .is_some()
    }
}

impl<'a> Allocator for BuddyHeapAllocator<'a> {

    /// Allocate a new block of size `size` on alignment `align`.
    ///
    /// # Arguments:
    /// + `size`: the amount of memory to allocate (in bytes)
    /// + `align`: the alignment for the allocation request
    ///
    /// # Returns:
    /// + `Some(*mut u8)` if the request was allocated successfully
    /// + `None` if the allocator is out of memory or if the request was
    ///     invalid.
    unsafe fn allocate( &mut self
                      , size: usize
                      , align: usize)
                      -> Option<*mut u8> {
        trace!(target: "alloc", "allocate() was called!");
        // First, compute the allocation order for this request
        self.alloc_order(size, align)
            .and_then(|order| if order > self.free_lists.len() - 1 { None }
                              else {Some(order)} )
            // If the allocation order is defined, then we try to allocate
            // a block of that order. Otherwise, the request is invalid.
            .and_then(|min_order| {
                trace!( target: "alloc"
                      , "in allocate(): min alloc order is {}", min_order);
                // Starting at the minimum possible order...
                // TODO: this is ugly and not FP, rewrite.
                for order in min_order..self.free_lists.len() {
                    if let Some(block) = self.pop_block(order) {
                        trace!(target: "alloc", "in allocate(): found block");
                        if order > min_order {
                            trace!( target: "alloc"
                                  , "in allocate(): order {} is less than \
                                     minimum ({}), splitting."
                                  , order, min_order);
                            self.split_block(block, order, min_order);
                            trace!( target: "alloc"
                                  , "in allocate(): split_block() done");

                        }
                        return Some(block)
                    }
                }
                None
            })
    }

    /// Release an allocated block of memory.
    ///
    /// The `size` and `align` parameters _must_ be the same as the original
    /// size and alignment of the frame being deallocated, otherwise our
    /// heap will become corrupted.
    ///
    /// # Arguments:
    /// + `frame`: a pointer to the block of memory to deallocate
    /// + `size`: the size of the block being deallocated
    /// + `align`: the alignment of the block being deallocated
    unsafe fn deallocate( &mut self
                        , block: *mut u8
                        , old_size: usize
                        , align: usize ) {
        let min_order = self.alloc_order(old_size, align)
                            .unwrap();

        // Check if the deallocated block's buddy block is also free.
        // If it is, merge the two blocks.
        let mut new_block = block;
        for order in min_order..self.free_lists.len() {
            // If there is a buddy for this block of the given order...
            if let Some(buddy) = self.get_buddy(order, block) {
                // ...and if the buddy was free...
                if self.remove_block(order, buddy) {
                    // ...merge the buddy with the new block (just use
                    // the lower address), and keep going.
                    new_block = min(new_block, buddy);
                    continue;
                }
            }
            // Otherwise, if we've run out of free buddies, push the new
            // merged block onto the free lsit and return.
            self.push_block(new_block, order);
            return;
        }
    }
}
<commit_msg>fix(alloc::buddy): Unique::get/get_mut were renamed to as_ref/as_mut<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Simple buddy-block allocator

mod math;
#[cfg(feature = "buddy_as_system")]
pub mod system;

use super::Framesque;
use self::math::PowersOf2;

use core::mem;
use core::cmp::{max, min};
use core::ptr::Unique;

use intrusive::list::{List, Node};
use intrusive::rawlink::RawLink;

use memory::alloc::Allocator;
use memory::arch::PAGE_SIZE;

#[cfg(target_os = "linux")]
#[cfg(test)]
mod test;

/// A `FreeList` is a list of unique free blocks
pub type FreeList = List<Unique<FreeBlock>, FreeBlock>;

/// A free block header stores a pointer to the next and previous free blocks.
pub struct FreeBlock { next: RawLink<FreeBlock>
                     , prev: RawLink<FreeBlock>
                     }

impl Framesque for FreeBlock {
    #[inline] fn as_ptr(&self) -> *mut u8 {
        unsafe { mem::transmute(self) } // HOPEFULLY this is good
    }
}

impl Node for FreeBlock {
    #[inline] fn prev(&self) -> &RawLink<FreeBlock> {
        &self.prev
    }
    #[inline] fn next(&self) -> &RawLink<FreeBlock> {
        &self.next
    }
    #[inline] fn prev_mut(&mut self) -> &mut RawLink<FreeBlock> {
        &mut self.prev
    }
    #[inline] fn next_mut(&mut self) -> &mut RawLink<FreeBlock> {
        &mut self.next
    }
}

// Variadic macro for taking the maximum of n > 2 numbers.
// because I'm lazy.
macro_rules! max {
    ($x:expr) => ($x);
    ($x:expr, $($xs:expr),+) => (max($x, max!($($xs),+)));
}

/// Structure with data for implementing the buddy block allocation strategy.
pub struct BuddyHeapAllocator<'a> {
    /// Address of the base of the heap. This must be aligned
    /// on a `MIN_ALIGN` boundary.
    pub start_addr: *mut u8
  , /// The allocator's free list
    free_lists: &'a mut [FreeList]
  , /// Number of blocks in the heap (must be a power of 2)
    pub heap_size: usize
  , /// Minimum block size
    pub min_block_size: usize
}

impl<'a> BuddyHeapAllocator<'a> {
    /// Construct a new `BuddyHeapAllocator`.
    ///
    /// # Arguments:
    /// + `start_addr`: a pointer to the start location of the heap
    /// + `free_lists`: an array of `FreeList`s. The cardinality
    ///     of the `free_lists` array should be equal to the maximum
    ///     allocateable order.
    /// + `heap_size`: the size of the heap (in bytes)
    ///
    /// # Returns:
    /// + A new `BuddyHeapAllocator`, obviously.
    pub unsafe fn new( start_addr: *mut u8
                     , free_lists: &'a mut [FreeList]
                     , heap_size: usize)
                     -> BuddyHeapAllocator<'a> {
        // Cache the number of free lists hopefully saving performance.
        let n_free_lists = free_lists.len();

        assert!( !start_addr.is_null()
                , "Heap start address cannot be null." );
        assert!( n_free_lists > 0
               , "Allocator must have at least one free list.");
        assert!( start_addr as usize & (PAGE_SIZE-1) as usize == 0
               , "Heap start address must be aligned on a 4k boundary.");

        let min_block_size = heap_size >> (n_free_lists - 1);

        assert!( heap_size >= min_block_size
               , "Heap must be large enough to contain at least one block.");
        assert!( min_block_size >= mem::size_of::<FreeBlock>()
               , "Minimum block size must be large enough to contain \
                  the free block header.");
        assert!( heap_size.is_pow2()
               , "Heap size must be a power of 2.");

    //    // We must have one free list per possible heap block size.
    //    assert_eq!(min_block_size *
    //               (2u32.pow(free_lists.len() as u32 - 1)) as usize,
    //               heap_size);

        // Zero out the free lists in case we were passed existing data.
        for list in free_lists.iter_mut() {
            *list = FreeList::new();
        }

        let mut heap
            = BuddyHeapAllocator { start_addr: start_addr
                                 , free_lists: free_lists
                                 , heap_size: heap_size
                                 , min_block_size: min_block_size
                                 };

        // the order needed to allocate the entire heap as a single block
        let root_order
            = heap.alloc_order(heap_size, 1)
                  .expect("Couldn't determine heap root allocation order!\
                           This should be (as far as I know) impossible.\
                           Something is seriously amiss.");

        // Push the entire heap onto the free lists as the first block.
        heap.push_block(start_addr, root_order);
        heap
    }

    /// Add a block of max order
    pub unsafe fn add_block(&mut self, block: *mut u8) {
        let order = self.free_lists.len() -1;
        self.push_block(block, order);
    }

    /// Computes the size of an allocation request.
    ///
    /// # Arguments
    /// + `size`: A `usize` containing the size of the request
    /// + `align`: A `usize` containing the alignment of the request
    ///
    /// # Returns
    /// + `None` if the request is invalid
    /// + `Some(usize)` containing the size needed if the request is valid
    #[inline]
    pub fn alloc_size(&self, size: usize, align: usize) -> Option<usize> {
        // Pre-check if this is a valid allocation request:
        //  - allocations must be aligned on power of 2 boundaries
        //  - we cannot allocate requests with alignments greater than the
        //    base alignment of the heap without jumping through a bunch of
        //    hoops.
        if !align.is_pow2() || align > PAGE_SIZE as usize {
            None
        // If the request is valid, compute the size we need to allocate
        } else {
            // the allocation size for the request is the next power of 2
            // after the size of the request, the alignment of the request,
            // or the minimum block size (whichever is greatest).
            let alloc_size = max!(size, self.min_block_size, align).next_pow2();

            if alloc_size > self.heap_size {
                // if the calculated size is greater than the size of the heap,
                // we (obviously) cannot allocate this request.
                None
            } else {
                // otherwise, return the calculated size.
                Some(alloc_size)
            }
        }
    }

    /// Computes the order of an allocation request.
    ///
    /// The "order" of an allocation refers to the number of times we need to
    /// double the minimum block size to get a large enough block for that
    /// allocation.
    #[inline]
    pub fn alloc_order(&self, size: usize, align: usize) -> Option<usize> {
        trace!(target: "alloc", "TRACE: alloc_order() called");
        self.alloc_size(size, align)
            .map(|s| {
                trace!(target: "alloc"
                      , "in alloc_order(): alloc_size() returned {}", s);
                s.log2() - self.min_block_size.log2()
            })
    }

    /// Computes the size  allocated for a request of the given order.
    #[inline]
    fn order_alloc_size(&self, order: usize) -> usize {
        1 << (self.min_block_size.log2() + order)
    }

    #[inline]
    unsafe fn push_block(&mut self, ptr: *mut u8, order: usize) {
        self.free_lists[order]
            .push_front(Unique::new(ptr as *mut FreeBlock))
    }

    #[inline]
    unsafe fn pop_block(&mut self, order: usize) -> Option<*mut u8>{
        self.free_lists[order]
            .pop_front()
            .map(|block| block.as_ref().as_ptr())
    }

    /// Splits a block
    unsafe fn split_block( &mut self
                         , block: *mut u8
                         , old_order: usize
                         , new_order: usize ) {
        trace!( target: "alloc"
              , "split_block() was called, target order: {}.", new_order);

        assert!( new_order < old_order
               , "Cannot split a block larger than its current order!");
        assert!( old_order <= self.free_lists.len()
               , "Order is too large for this allocator!");

        let mut split_size = self.order_alloc_size(old_order);
        // let mut curr_order = order;
        for order in (new_order..old_order).rev() {
            split_size >>= 1;

            // let split_size = self.order_alloc_size(order);
            self.push_block(block.offset(split_size as isize), order);

            trace!( target: "alloc"
                  , "split block successfully, order: {}, split size: {}"
                  , order, split_size );

        }
    }

    /// Finds the buddy block for a given block.
    ///
    /// # Arguments
    /// + `order`: the order of the block to find a buddy for
    /// + `block`: a pointer to the block to find a buddy for
    ///
    /// # Returns
    /// + `Some(*mut u8)` pointing to the buddy block if a buddy was found
    /// + `None` if the block was the size of the entire heap
    pub unsafe fn get_buddy( &self
                           , order: usize
                           , block: *mut u8)
                            -> Option<*mut u8> {
        // Determine the size of the block allocated for the given order
        let block_size = self.order_alloc_size(order);
        if block_size < self.heap_size {
            // Determine the block's position in the heap.
            let block_pos = (block as usize) - (self.start_addr as usize);
            // Calculate the block's buddy by XORing the block's position
            // in the heap with its size.
            Some(self.start_addr.offset((block_pos ^ block_size) as isize))
        } else {
            // If the block is the size of the entire heap, it (obviously)
            // cannot have a buddy block.
            None
        }
    }

    /// Finds and removes the target block from the free list.
    ///
    /// # Arguments
    /// + `order`: the order of the free list to remove the block from_raw
    /// + `block`: a pointer to the block to remove
    ///
    /// # Returns
    /// + `true` if the block was found and removed from the free List
    /// + `false` if the block was not found
    pub fn remove_block(&mut self, order: usize, block: *mut u8) -> bool {
        self.free_lists[order]
            .cursor_mut()
            .find_and_remove(|b| b as *const FreeBlock as *const u8 == block)
            .is_some()
    }
}

impl<'a> Allocator for BuddyHeapAllocator<'a> {

    /// Allocate a new block of size `size` on alignment `align`.
    ///
    /// # Arguments:
    /// + `size`: the amount of memory to allocate (in bytes)
    /// + `align`: the alignment for the allocation request
    ///
    /// # Returns:
    /// + `Some(*mut u8)` if the request was allocated successfully
    /// + `None` if the allocator is out of memory or if the request was
    ///     invalid.
    unsafe fn allocate( &mut self
                      , size: usize
                      , align: usize)
                      -> Option<*mut u8> {
        trace!(target: "alloc", "allocate() was called!");
        // First, compute the allocation order for this request
        self.alloc_order(size, align)
            .and_then(|order| if order > self.free_lists.len() - 1 { None }
                              else {Some(order)} )
            // If the allocation order is defined, then we try to allocate
            // a block of that order. Otherwise, the request is invalid.
            .and_then(|min_order| {
                trace!( target: "alloc"
                      , "in allocate(): min alloc order is {}", min_order);
                // Starting at the minimum possible order...
                // TODO: this is ugly and not FP, rewrite.
                for order in min_order..self.free_lists.len() {
                    if let Some(block) = self.pop_block(order) {
                        trace!(target: "alloc", "in allocate(): found block");
                        if order > min_order {
                            trace!( target: "alloc"
                                  , "in allocate(): order {} is less than \
                                     minimum ({}), splitting."
                                  , order, min_order);
                            self.split_block(block, order, min_order);
                            trace!( target: "alloc"
                                  , "in allocate(): split_block() done");

                        }
                        return Some(block)
                    }
                }
                None
            })
    }

    /// Release an allocated block of memory.
    ///
    /// The `size` and `align` parameters _must_ be the same as the original
    /// size and alignment of the frame being deallocated, otherwise our
    /// heap will become corrupted.
    ///
    /// # Arguments:
    /// + `frame`: a pointer to the block of memory to deallocate
    /// + `size`: the size of the block being deallocated
    /// + `align`: the alignment of the block being deallocated
    unsafe fn deallocate( &mut self
                        , block: *mut u8
                        , old_size: usize
                        , align: usize ) {
        let min_order = self.alloc_order(old_size, align)
                            .unwrap();

        // Check if the deallocated block's buddy block is also free.
        // If it is, merge the two blocks.
        let mut new_block = block;
        for order in min_order..self.free_lists.len() {
            // If there is a buddy for this block of the given order...
            if let Some(buddy) = self.get_buddy(order, block) {
                // ...and if the buddy was free...
                if self.remove_block(order, buddy) {
                    // ...merge the buddy with the new block (just use
                    // the lower address), and keep going.
                    new_block = min(new_block, buddy);
                    continue;
                }
            }
            // Otherwise, if we've run out of free buddies, push the new
            // merged block onto the free lsit and return.
            self.push_block(new_block, order);
            return;
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Alexander Reece
//
// 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.

#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]
#![cfg_attr(not(feature="clippy"), allow(unknown_lints))]

extern crate phf;

use phf::OrderedMap;
use std::hash::{Hash, Hasher};

#[cfg(feature = "amqp0-build-specs")]
include!(concat!(env!("OUT_DIR"), "/amqp0.rs"));
#[cfg(not(feature = "amqp0-build-specs"))]
include!(concat!("lib.pregen.rs"));

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Version {
    minor: u16,
    revision: u16,
}

#[derive(Debug, Clone)]
pub struct Spec {
    name: &'static str,
    classes: &'static OrderedMap<&'static str, Class>,
    constants: &'static OrderedMap<&'static str, Constant>,
    domains: &'static OrderedMap<&'static str, &'static str>,
    frame_types: &'static OrderedMap<&'static str, Constant>,
    response_codes: &'static OrderedMap<&'static str, Constant>,
    version: Version,
}

use phf::PhfHash;

impl PartialEq for Spec {
    fn eq(&self, other: &Spec) -> bool {

        fn is_map_eq<K, V>(map1: &OrderedMap<K, V>, map2: &OrderedMap<K, V>) -> bool
            where K: Eq + PhfHash,
                  V: Eq
        {
            map1.entries().all(|(k, v1)| {
                match map2.get(k) {
                    Some(v2) if v1 == v2 => true,
                    _ => false,
                }
            })
        }

        self.name == other.name
            && self.version == other.version
            && is_map_eq(self.classes, other.classes)
            && is_map_eq(self.constants, other.constants)
            && is_map_eq(self.domains, other.domains)
            && is_map_eq(self.frame_types, other.frame_types)
            && is_map_eq(self.response_codes, other.response_codes)
    }
}
impl Eq for Spec {}

impl Hash for Spec {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        for (name, class) in self.classes {
            name.hash(state);
            class.hash(state);
        }
        for (name, value) in self.constants {
            name.hash(state);
            value.hash(state);
        }
        for (name, mapping) in self.domains {
            name.hash(state);
            mapping.hash(state);
        }
        for (name, value) in self.frame_types {
            name.hash(state);
            value.hash(state);
        }
        for (name, response_code) in self.response_codes {
            name.hash(state);
            response_code.hash(state);
        }
        self.version.hash(state)
    }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Class {
    name: &'static str,
    fields: &'static [ClassField],
    index: u16,
    methods: &'static [ClassMethod],
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct ClassMethod {
    name: &'static str,
    index: u16,
    response: Option<&'static str>,
    fields: &'static [ClassMethodField],
    is_synchronous: bool,

    chassis_server: Option<&'static str>,
    chassis_client: Option<&'static str>,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct ClassField {
    name: &'static str,
    domain: &'static str,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct ClassMethodField {
    name: &'static str,
    domain: &'static str,
    assertions: &'static [ClassMethodFieldAssertion],
    is_reserved: bool,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum ClassMethodFieldAssertion {
    Null,
    NotNull,
    ChannelMax,
    NotZero,
    Enum(&'static [&'static str]),
    Length(usize),
    Regexp(&'static str),
    Syntax(&'static str),
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Constant {
    name: &'static str,
    value: u32,
    class: Option<&'static str>,
}

impl Spec {
    pub fn name(&self) -> &'static str {
        self.name
    }
    pub fn classes(&self) -> &'static OrderedMap<&'static str, Class> {
        self.classes
    }
    pub fn domains(&self) -> &'static OrderedMap<&'static str, &'static str> {
        self.domains
    }
    pub fn frame_types(&self) -> &'static OrderedMap<&'static str, Constant> {
        self.frame_types
    }
    pub fn response_codes(&self) -> &'static OrderedMap<&'static str, Constant> {
        self.response_codes
    }
    pub fn version(&self) -> &Version {
        &self.version
    }
}

impl Class {
    pub fn name(&self) -> &'static str {
        self.name
    }

    pub fn fields(&self) -> &'static [ClassField] {
        &self.fields
    }

    pub fn index(&self) -> u16 {
        self.index
    }

    pub fn methods(&self) -> &'static [ClassMethod] {
        &self.methods
    }
}

impl ClassField {
    pub fn name(&self) -> &'static str {
        &self.name
    }

    pub fn domain(&self) -> &'static str {
        self.domain
    }
}

impl ClassMethod {
    pub fn index(&self) -> u16 {
        self.index
    }
    pub fn name(&self) -> &str {
        &self.name
    }
    pub fn fields(&self) -> &'static [ClassMethodField] {
        &self.fields
    }
}

impl ClassMethodField {
    pub fn name(&self) -> &'static str {
        &self.name
    }

    pub fn domain(&self) -> &'static str {
        self.domain
    }

    pub fn assertions(&self) -> &'static [ClassMethodFieldAssertion] {
        &self.assertions
    }

    pub fn is_reserved(&self) -> bool {
        self.is_reserved
    }
}

impl Constant {
    pub fn name(&self) -> &'static str {
        self.name
    }

    pub fn value(&self) -> u32 {
        self.value
    }

    pub fn class(&self) -> Option<&'static str> {
        self.class
    }
}

impl Version {
    pub fn minor(&self) -> u16 {
        self.minor
    }

    pub fn revision(&self) -> u16 {
        self.revision
    }
}<commit_msg>[specs] Remove redundant references<commit_after>// Copyright 2016 Alexander Reece
//
// 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.

#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]
#![cfg_attr(not(feature="clippy"), allow(unknown_lints))]

extern crate phf;

use phf::OrderedMap;
use std::hash::{Hash, Hasher};

#[cfg(feature = "amqp0-build-specs")]
include!(concat!(env!("OUT_DIR"), "/amqp0.rs"));
#[cfg(not(feature = "amqp0-build-specs"))]
include!(concat!("lib.pregen.rs"));

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Version {
    minor: u16,
    revision: u16,
}

#[derive(Debug, Clone)]
pub struct Spec {
    name: &'static str,
    classes: &'static OrderedMap<&'static str, Class>,
    constants: &'static OrderedMap<&'static str, Constant>,
    domains: &'static OrderedMap<&'static str, &'static str>,
    frame_types: &'static OrderedMap<&'static str, Constant>,
    response_codes: &'static OrderedMap<&'static str, Constant>,
    version: Version,
}

use phf::PhfHash;

impl PartialEq for Spec {
    fn eq(&self, other: &Spec) -> bool {

        fn is_map_eq<K, V>(map1: &OrderedMap<K, V>, map2: &OrderedMap<K, V>) -> bool
            where K: Eq + PhfHash,
                  V: Eq
        {
            map1.entries().all(|(k, v1)| {
                match map2.get(k) {
                    Some(v2) if v1 == v2 => true,
                    _ => false,
                }
            })
        }

        self.name == other.name
            && self.version == other.version
            && is_map_eq(self.classes, other.classes)
            && is_map_eq(self.constants, other.constants)
            && is_map_eq(self.domains, other.domains)
            && is_map_eq(self.frame_types, other.frame_types)
            && is_map_eq(self.response_codes, other.response_codes)
    }
}
impl Eq for Spec {}

impl Hash for Spec {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        for (name, class) in self.classes {
            name.hash(state);
            class.hash(state);
        }
        for (name, value) in self.constants {
            name.hash(state);
            value.hash(state);
        }
        for (name, mapping) in self.domains {
            name.hash(state);
            mapping.hash(state);
        }
        for (name, value) in self.frame_types {
            name.hash(state);
            value.hash(state);
        }
        for (name, response_code) in self.response_codes {
            name.hash(state);
            response_code.hash(state);
        }
        self.version.hash(state)
    }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Class {
    name: &'static str,
    fields: &'static [ClassField],
    index: u16,
    methods: &'static [ClassMethod],
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct ClassMethod {
    name: &'static str,
    index: u16,
    response: Option<&'static str>,
    fields: &'static [ClassMethodField],
    is_synchronous: bool,

    chassis_server: Option<&'static str>,
    chassis_client: Option<&'static str>,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct ClassField {
    name: &'static str,
    domain: &'static str,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct ClassMethodField {
    name: &'static str,
    domain: &'static str,
    assertions: &'static [ClassMethodFieldAssertion],
    is_reserved: bool,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum ClassMethodFieldAssertion {
    Null,
    NotNull,
    ChannelMax,
    NotZero,
    Enum(&'static [&'static str]),
    Length(usize),
    Regexp(&'static str),
    Syntax(&'static str),
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Constant {
    name: &'static str,
    value: u32,
    class: Option<&'static str>,
}

impl Spec {
    pub fn name(&self) -> &'static str {
        self.name
    }
    pub fn classes(&self) -> &'static OrderedMap<&'static str, Class> {
        self.classes
    }
    pub fn domains(&self) -> &'static OrderedMap<&'static str, &'static str> {
        self.domains
    }
    pub fn frame_types(&self) -> &'static OrderedMap<&'static str, Constant> {
        self.frame_types
    }
    pub fn response_codes(&self) -> &'static OrderedMap<&'static str, Constant> {
        self.response_codes
    }
    pub fn version(&self) -> &Version {
        &self.version
    }
}

impl Class {
    pub fn name(&self) -> &'static str {
        self.name
    }

    pub fn fields(&self) -> &'static [ClassField] {
        self.fields
    }

    pub fn index(&self) -> u16 {
        self.index
    }

    pub fn methods(&self) -> &'static [ClassMethod] {
        self.methods
    }
}

impl ClassField {
    pub fn name(&self) -> &'static str {
        self.name
    }

    pub fn domain(&self) -> &'static str {
        self.domain
    }
}

impl ClassMethod {
    pub fn index(&self) -> u16 {
        self.index
    }
    pub fn name(&self) -> &str {
        self.name
    }
    pub fn fields(&self) -> &'static [ClassMethodField] {
        self.fields
    }
}

impl ClassMethodField {
    pub fn name(&self) -> &'static str {
        self.name
    }

    pub fn domain(&self) -> &'static str {
        self.domain
    }

    pub fn assertions(&self) -> &'static [ClassMethodFieldAssertion] {
        self.assertions
    }

    pub fn is_reserved(&self) -> bool {
        self.is_reserved
    }
}

impl Constant {
    pub fn name(&self) -> &'static str {
        self.name
    }

    pub fn value(&self) -> u32 {
        self.value
    }

    pub fn class(&self) -> Option<&'static str> {
        self.class
    }
}

impl Version {
    pub fn minor(&self) -> u16 {
        self.minor
    }

    pub fn revision(&self) -> u16 {
        self.revision
    }
}<|endoftext|>


<commit_before>extern crate syn;

#[macro_use]
extern crate quote;

use std::env;
use std::path::{PathBuf, Path as FilePath};

use syn::*;
use quote::Tokens;

struct Context {
    location_ident: Ident,
    location_expr: Expr,
    location_type: Ty,
    deref_self_expr: Expr,
}

impl Context {
    pub fn new() -> Self {
        Context {
            location_ident: Ident::from("location"),

            location_expr: Expr::from(ExprKind::Path(None, Path::from("location"))),

            location_type: Ty::Path(None, Path {
                global: false,
                segments: vec![PathSegment {
                    ident: Ident::from("Option"),
                    parameters: PathParameters::AngleBracketed(AngleBracketedParameterData {
                        lifetimes: vec![],
                        types: vec![Ty::Path(None, Path::from("Span"))],
                        bindings: vec![]
                    })
                }]
            }),

            deref_self_expr: Expr::from(ExprKind::Unary(
                UnOp::Deref,
                Box::new(Expr::from(ExprKind::Path(None, Path::from("self"))))
            ))
        }
    }

    fn expand_tracking_ref_data(&self, path: Path, data: &VariantData, mutability: Mutability) -> Arm {
        let location_pat = Pat::Ident(BindingMode::ByRef(mutability), self.location_ident.clone(), None);

        let (pat, field) = match *data {
            VariantData::Struct(_) => (
                Pat::Struct(path, vec![FieldPat {
                    ident: self.location_ident.clone(),
                    pat: Box::new(location_pat),
                    is_shorthand: true
                }], true),
                data.fields().iter().find(|field| field.ident.as_ref() == Some(&self.location_ident)).unwrap()
            ),
            VariantData::Tuple(ref fields) => (
                Pat::TupleStruct(path, {
                    let mut v = vec![location_pat];
                    v.extend(std::iter::repeat(Pat::Wild).take(fields.len() - 1));
                    v
                }, None),
                &data.fields()[0]
            ),
            VariantData::Unit => panic!("Empty unit is not trackable")
        };

        let expr = self.location_expr.clone();

        Arm {
            attrs: vec![],
            pats: vec![pat],
            guard: None,
            body: Box::new(if field.ty == self.location_type {
                expr
            } else {
                Expr::from(ExprKind::MethodCall(
                    Ident::from(if mutability == Mutability::Immutable { "tracking_ref" } else { "tracking_mut" }),
                    vec![],
                    vec![expr]
                ))
            })
        }
    }

    pub fn expand_tracking_ref(&self, ast: &MacroInput, mutability: Mutability) -> Tokens {
        // Helper is provided for handling complex generic types correctly and effortlessly
        let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();

        let (impl_name, method) = if mutability == Mutability::Immutable {
            (Token::Ident(Ident::from("TrackingRef")), quote! {
                tracking_ref(&self) -> &Option
            })
        } else {
            (Token::Ident(Ident::from("TrackingMut")), quote! {
                tracking_mut(&mut self) -> &mut Option
            })
        };

        // Used in the quasi-quotation below as `#name`
        let name = &ast.ident;

        let body = Expr::from(ExprKind::Match(
            Box::new(self.deref_self_expr.clone()),
            match ast.body {
                Body::Struct(ref data) => {
                    vec![self.expand_tracking_ref_data(Path::from(name.clone()), data, mutability)]
                },
                Body::Enum(ref variants) => {
                    variants.iter().map(|var| {
                        let path = Path {
                            global: false,
                            segments: vec![
                                PathSegment::from(name.clone()),
                                PathSegment::from(var.ident.clone())
                            ]
                        };
                        self.expand_tracking_ref_data(path, &var.data, mutability)
                    }).collect()
                }
            }
        ));

        quote! {
            // The generated impl
            impl #impl_generics #impl_name for #name #ty_generics #where_clause {
                fn #method {
                    #body
                }
            }
        }
    }

    fn expand_untrack_data(&self, path: Path, data: &VariantData) -> Arm {
        let (pat, idents) = match *data {
            VariantData::Struct(ref fields) => {
                let field_pats: Vec<_> = fields.iter().map(|field| {
                    let ident = field.ident.clone().unwrap();

                    FieldPat {
                        ident: ident.clone(),
                        pat: Box::new(Pat::Ident(BindingMode::ByRef(Mutability::Mutable), ident, None)),
                        is_shorthand: true
                    }
                }).collect();

                let idents: Vec<_> = fields.iter().map(|field| field.ident.clone().unwrap()).collect();

                (
                    Pat::Struct(path, field_pats, false),
                    idents
                )
            },
            VariantData::Tuple(ref fields) => {
                let pats: Vec<_> = fields.iter().enumerate().map(|(i, _)| {
                    Pat::Ident(BindingMode::ByRef(Mutability::Mutable), Ident::from(format!("_f{}", i)), None)
                }).collect();

                let idents: Vec<_> = fields.iter().enumerate().map(|(i, _)| {
                    Ident::from(format!("_f{}", i))
                }).collect();

                (
                    Pat::TupleStruct(path, pats, None),
                    idents
                )
            },
            VariantData::Unit => {
                (
                    Pat::Path(None, path),
                    vec![]
                )
            }
        };

        let expr = Expr::from(ExprKind::Block(BlockCheckMode::Default, Block {
            stmts: idents.into_iter().map(|ident| {
                Stmt::Semi(Box::new(Expr::from(ExprKind::MethodCall(
                    Ident::from("untrack"),
                    vec![],
                    vec![Expr::from(ExprKind::Path(None, Path::from(ident)))]
                ))))
            }).collect()
        }));

        Arm {
            attrs: vec![],
            pats: vec![pat],
            guard: None,
            body: Box::new(expr)
        }
    }

    pub fn expand_untrack(&self, ast: &MacroInput) -> Tokens {
        let mut generics = ast.generics.clone();

        let bound = TyParamBound::Trait(PolyTraitRef {
            bound_lifetimes: vec![],
            trait_ref: Path::from("Untrack")
        }, TraitBoundModifier::None);

        for ty in &mut generics.ty_params {
            ty.bounds.push(bound.clone());
        }

        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

        let name = &ast.ident;

        let body = Expr::from(ExprKind::Match(
            Box::new(self.deref_self_expr.clone()),
            match ast.body {
                Body::Struct(ref data) => {
                    vec![self.expand_untrack_data(Path::from(name.clone()), data)]
                },
                Body::Enum(ref variants) => {
                    variants.iter().map(|var| {
                        let path = Path {
                            global: false,
                            segments: vec![
                                PathSegment::from(name.clone()),
                                PathSegment::from(var.ident.clone())
                            ]
                        };
                        self.expand_untrack_data(path, &var.data)
                    }).collect()
                }
            }
        ));

        quote! {
            // The generated impl
            impl #impl_generics Untrack for #name #ty_generics #where_clause {
                fn untrack(&mut self) {
                    #body
                }
            }
        }
    }
}

fn main() {
    let out_dir = env::var_os("OUT_DIR").unwrap();
    let out_path = FilePath::new(&out_dir);

    let src_path = {
        let mut buf = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        buf.push("src");
        buf
    };

    let context = Context::new();

    let registry = {
        let mut registry = Registry::new();

        registry.add_derive("TrackingRef", |input: MacroInput| {
            let tokens = context.expand_tracking_ref(&input, Mutability::Immutable);
            Ok(Expanded {
                new_items: parse_items(&tokens.to_string())?,
                original: Some(input),
            })
        });

        registry.add_derive("TrackingMut", |input: MacroInput| {
            let tokens = context.expand_tracking_ref(&input, Mutability::Mutable);
            Ok(Expanded {
                new_items: parse_items(&tokens.to_string())?,
                original: Some(input),
            })
        });

        registry.add_derive("Untrack", |input: MacroInput| {
            let tokens = context.expand_untrack(&input);
            Ok(Expanded {
                new_items: parse_items(&tokens.to_string())?,
                original: Some(input),
            })
        });

        registry
    };

    for entry in src_path.read_dir().unwrap() {
        let path = entry.unwrap().path();
        let file_name = path.file_name().unwrap();
        if file_name != "lib.rs" {
            let dest = out_path.join(file_name);
            registry.expand_file(&path, &dest).unwrap();
        }
    }
}<commit_msg>Simplify Untrack expander<commit_after>extern crate syn;

#[macro_use]
extern crate quote;

use std::env;
use std::path::{PathBuf, Path as FilePath};

use syn::*;
use quote::Tokens;

struct Context {
    location_ident: Ident,
    location_expr: Expr,
    location_type: Ty,
    deref_self_expr: Expr,
}

impl Context {
    pub fn new() -> Self {
        Context {
            location_ident: Ident::from("location"),

            location_expr: Expr::from(ExprKind::Path(None, Path::from("location"))),

            location_type: Ty::Path(None, Path {
                global: false,
                segments: vec![PathSegment {
                    ident: Ident::from("Option"),
                    parameters: PathParameters::AngleBracketed(AngleBracketedParameterData {
                        lifetimes: vec![],
                        types: vec![Ty::Path(None, Path::from("Span"))],
                        bindings: vec![]
                    })
                }]
            }),

            deref_self_expr: Expr::from(ExprKind::Unary(
                UnOp::Deref,
                Box::new(Expr::from(ExprKind::Path(None, Path::from("self"))))
            ))
        }
    }

    fn expand_tracking_ref_data(&self, path: Path, data: &VariantData, mutability: Mutability) -> Arm {
        let location_pat = Pat::Ident(BindingMode::ByRef(mutability), self.location_ident.clone(), None);

        let (pat, field) = match *data {
            VariantData::Struct(_) => (
                Pat::Struct(path, vec![FieldPat {
                    ident: self.location_ident.clone(),
                    pat: Box::new(location_pat),
                    is_shorthand: true
                }], true),
                data.fields().iter().find(|field| field.ident.as_ref() == Some(&self.location_ident)).unwrap()
            ),
            VariantData::Tuple(ref fields) => (
                Pat::TupleStruct(path, {
                    let mut v = Vec::with_capacity(fields.len());
                    v.push(location_pat);
                    v.extend(std::iter::repeat(Pat::Wild).take(fields.len() - 1));
                    v
                }, None),
                &data.fields()[0]
            ),
            VariantData::Unit => panic!("Empty unit is not trackable")
        };

        let expr = self.location_expr.clone();

        Arm {
            attrs: vec![],
            pats: vec![pat],
            guard: None,
            body: Box::new(if field.ty == self.location_type {
                expr
            } else {
                Expr::from(ExprKind::MethodCall(
                    Ident::from(if mutability == Mutability::Immutable { "tracking_ref" } else { "tracking_mut" }),
                    vec![],
                    vec![expr]
                ))
            })
        }
    }

    pub fn expand_tracking_ref(&self, ast: &MacroInput, mutability: Mutability) -> Tokens {
        // Helper is provided for handling complex generic types correctly and effortlessly
        let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();

        let (impl_name, method) = if mutability == Mutability::Immutable {
            (Ident::from("TrackingRef"), quote! {
                tracking_ref(&self) -> &Option
            })
        } else {
            (Ident::from("TrackingMut"), quote! {
                tracking_mut(&mut self) -> &mut Option
            })
        };

        // Used in the quasi-quotation below as `#name`
        let name = &ast.ident;

        let body = Expr::from(ExprKind::Match(
            Box::new(self.deref_self_expr.clone()),
            match ast.body {
                Body::Struct(ref data) => {
                    vec![self.expand_tracking_ref_data(Path::from(name.clone()), data, mutability)]
                },
                Body::Enum(ref variants) => {
                    variants.iter().map(|var| {
                        let path = Path {
                            global: false,
                            segments: vec![
                                PathSegment::from(name.clone()),
                                PathSegment::from(var.ident.clone())
                            ]
                        };
                        self.expand_tracking_ref_data(path, &var.data, mutability)
                    }).collect()
                }
            }
        ));

        quote! {
            // The generated impl
            impl #impl_generics #impl_name for #name #ty_generics #where_clause {
                fn #method {
                    #body
                }
            }
        }
    }

    fn expand_untrack_data(&self, path: Path, data: &VariantData) -> Arm {
        let (pat, idents) = match *data {
            VariantData::Struct(ref fields) => {
                let mut field_pats = Vec::with_capacity(fields.len());
                let mut idents = Vec::with_capacity(fields.len());

                for field in fields {
                    let ident = field.ident.as_ref().unwrap();

                    field_pats.push(FieldPat {
                        ident: ident.clone(),
                        pat: Box::new(Pat::Ident(BindingMode::ByRef(Mutability::Mutable), ident.clone(), None)),
                        is_shorthand: true
                    });

                    idents.push(ident.clone());
                }

                (Pat::Struct(path, field_pats, false), idents)
            },
            VariantData::Tuple(ref fields) => {
                let mut pats = Vec::with_capacity(fields.len());
                let mut idents = Vec::with_capacity(fields.len());

                for i in 0..fields.len() {
                    let ident = Ident::from(format!("_f{}", i));

                    pats.push(Pat::Ident(BindingMode::ByRef(Mutability::Mutable), ident.clone(), None));

                    idents.push(ident);
                }

                (Pat::TupleStruct(path, pats, None), idents)
            },
            VariantData::Unit => {
                (Pat::Path(None, path), vec![])
            }
        };

        let expr = Expr::from(ExprKind::Block(BlockCheckMode::Default, Block {
            stmts: idents.into_iter().map(|ident| {
                Stmt::Semi(Box::new(Expr::from(ExprKind::MethodCall(
                    Ident::from("untrack"),
                    vec![],
                    vec![Expr::from(ExprKind::Path(None, Path::from(ident)))]
                ))))
            }).collect()
        }));

        Arm {
            attrs: vec![],
            pats: vec![pat],
            guard: None,
            body: Box::new(expr)
        }
    }

    pub fn expand_untrack(&self, ast: &MacroInput) -> Tokens {
        let mut generics = ast.generics.clone();

        let bound = TyParamBound::Trait(PolyTraitRef {
            bound_lifetimes: vec![],
            trait_ref: Path::from("Untrack")
        }, TraitBoundModifier::None);

        for ty in &mut generics.ty_params {
            ty.bounds.push(bound.clone());
        }

        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

        let name = &ast.ident;

        let body = Expr::from(ExprKind::Match(
            Box::new(self.deref_self_expr.clone()),
            match ast.body {
                Body::Struct(ref data) => {
                    vec![self.expand_untrack_data(Path::from(name.clone()), data)]
                },
                Body::Enum(ref variants) => {
                    variants.iter().map(|var| {
                        let path = Path {
                            global: false,
                            segments: vec![
                                PathSegment::from(name.clone()),
                                PathSegment::from(var.ident.clone())
                            ]
                        };
                        self.expand_untrack_data(path, &var.data)
                    }).collect()
                }
            }
        ));

        quote! {
            // The generated impl
            impl #impl_generics Untrack for #name #ty_generics #where_clause {
                fn untrack(&mut self) {
                    #body
                }
            }
        }
    }
}

fn main() {
    let out_dir = env::var_os("OUT_DIR").unwrap();
    let out_path = FilePath::new(&out_dir);

    let src_path = {
        let mut buf = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        buf.push("src");
        buf
    };

    let context = Context::new();

    let registry = {
        let mut registry = Registry::new();

        registry.add_derive("TrackingRef", |input| {
            let tokens = context.expand_tracking_ref(&input, Mutability::Immutable);
            Ok(Expanded {
                new_items: parse_items(&tokens.to_string())?,
                original: Some(input),
            })
        });

        registry.add_derive("TrackingMut", |input| {
            let tokens = context.expand_tracking_ref(&input, Mutability::Mutable);
            Ok(Expanded {
                new_items: parse_items(&tokens.to_string())?,
                original: Some(input),
            })
        });

        registry.add_derive("Untrack", |input| {
            let tokens = context.expand_untrack(&input);
            Ok(Expanded {
                new_items: parse_items(&tokens.to_string())?,
                original: Some(input),
            })
        });

        registry
    };

    for entry in src_path.read_dir().unwrap() {
        let path = entry.unwrap().path();
        let file_name = path.file_name().unwrap();
        if file_name != "lib.rs" {
            let dest = out_path.join(file_name);
            registry.expand_file(&path, &dest).unwrap();
        }
    }
}<|endoftext|>


<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! The Resolver is responsible for performing recursive queries to lookup domain names.
//!
//! This is a 100% in process DNS resolver. It *does not* use the Host OS' resolver. If what is desired is to use the Host OS' resolver, generally in the system's libc, then the `std::net::ToSocketAddrs` variant over `&str` should be used.
//!
//! Unlike the `trust-dns` client, this tries to provide a simpler interface to perform DNS queries. For update options, i.e. Dynamic DNS, the `trust-dns` crate must be used directly. The Resolver library is capable of searching multiple domains (this can be disabled by using an FQDN during lookup), dual-stack IPv4/IPv6 lookups, performing chained CNAME lookups, and features connection metric tracking for attempting to pick the best upstream DNS resolver.
//!
//! There are two types for performing DNS queries, `Resolver` and `ResolverFuture`. `Resolver` is the easiest to work with, it is a wrapper around `ResolverFuture`. `ResolverFuture` is a `Tokio` based async resolver, and can be used inside any `Tokio` based system.
//!
//! This as best as possible attempts to abide by the the DNS RFCs, please file issues at https://github.com/bluejekyll/trust-dns .
//!
//! # Usage
//!
//! ## Declare dependency
//!
//! ```toml
//! [dependency]
//! trust-dns-resolver = "^0.5"
//! ```
//!
//! ## Extern the crate for usage in the library
//!
//! ```rust
//! extern crate trust_dns_resolver;
//! ```
//!
//! ## Using the Synchronous Resolver
//!
//! This uses the default configuration, which sets the [Google Public DNS](https://developers.google.com/speed/public-dns/) as the upstream resolvers. Please see their [privacy statement](https://developers.google.com/speed/public-dns/privacy) for important information about what they track, many ISP's track similar information in DNS.
//!
//! ```rust
//! use std::net::*;
//! use trust_dns_resolver::Resolver;
//! use trust_dns_resolver::config::*;
//!
//! // Construct a new Resolver with default configuration options
//! let resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();
//!
//! // Lookup the IP addresses associated with a name.
//! // The final dot forces this to be an FQDN, otherwise the search rules as specified
//! //  in `ResolverOpts` will take effect. FQDN's are generally cheaper queries.
//! let response = resolver.lookup_ip("www.example.com.").unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.iter().next().expect("no addresses returned!");
//! if address.is_ipv4() {
//!     assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! } else {
//!     assert_eq!(address, IpAddr::V6(Ipv6Addr::new(0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946)));
//! }
//! ```
//!
//! ## Using the host system config
//!
//! On Unix systems, the `/etc/resolv.conf` can be used for configuration. Not all options specified in the host systems `resolv.conf` are applicable or compatible with this software. In addition there may be additional options supported which the host system does not. Example:
//!
//! ```rust,no_run
//! # use std::net::*;
//! # use trust_dns_resolver::Resolver;
//! // Use the host OS'es `/etc/resolv.conf`
//! # #[cfg(unix)]
//! let resolver = Resolver::from_system_conf().unwrap();
//! # #[cfg(unix)]
//! let response = resolver.lookup_ip("www.example.com.").unwrap();
//! ```
//!
//! ## Using the Tokio/Async Resolver
//!
//! For more advanced asynchronous usage, the ResolverFuture is integrated with Tokio. In fact, the ResolverFuture is used by the synchronous Resolver for all lookups.
//!
//! ```rust
//! # extern crate futures;
//! # extern crate tokio_core;
//! # extern crate trust_dns_resolver;
//! # fn main() {
//! use std::net::*;
//! use tokio_core::reactor::Core;
//! use trust_dns_resolver::ResolverFuture;
//! use trust_dns_resolver::config::*;
//!
//! // We need a Tokio reactor::Core to run the resolver
//! //  this is responsible for running all Future tasks and registering interest in IO channels
//! let mut io_loop = Core::new().unwrap();
//!
//! // Construct a new Resolver with default configuration options
//! let resolver = ResolverFuture::new(ResolverConfig::default(), ResolverOpts::default(), &io_loop.handle());
//!
//! // Lookup the IP addresses associated with a name.
//! // This returns a future that will lookup the IP addresses, it must be run in the Core to
//! //  to get the actual result.
//! let lookup_future = resolver.lookup_ip("www.example.com.");
//!
//! // Run the lookup until it resolves or errors
//! let mut response = io_loop.run(lookup_future).unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.iter().next().expect("no addresses returned!");
//! if address.is_ipv4() {
//!     assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! } else {
//!     assert_eq!(address, IpAddr::V6(Ipv6Addr::new(0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946)));
//! }
//! # }
//! ```
//!
//! Generally after a lookup in an asynchornous context, there would probably be a connection made to a server, for example:
//!
//! ```c
//! let result = io_loop.run(lookup_future.and_then(|ips| {
//!                                  let ip = ips.next().unwrap();
//!                                  TcpStream::connect()
//!                              }).and_then(|conn| /* do something with the connection... */)
//!                          ).unwrap();
//! ```
//!
//! It's beyond the scope of these examples to show how to deal with connection failures and looping etc. But if you wanted to say try a different address from the result set after a connection failure, it will be necessary to create a type that implements the `Future` trait. Inside the `Future::poll` method would be the place to implement a loop over the different IP addresses.

#![deny(missing_docs)]

#[macro_use]
extern crate error_chain;
extern crate futures;
extern crate lalrpop_util;
#[macro_use]
extern crate log;
extern crate lru_cache;
extern crate tokio_core;
extern crate trust_dns_proto;
#[cfg(all(feature = "ipconfig", target_os = "windows", target_pointer_width = "64"))]
extern crate ipconfig;

pub mod config;
pub mod error;
pub mod lookup_ip;
pub mod lookup;
pub mod lookup_state;
#[doc(hidden)]
pub mod name_server_pool;
mod resolver;
pub mod system_conf;
mod resolver_future;
mod hosts;

pub use resolver::Resolver;
pub use resolver_future::ResolverFuture;
pub use hosts::Hosts;

/// returns a version as specified in Cargo.toml
pub fn version() -> &'static str {
    env!("CARGO_PKG_VERSION")
}
<commit_msg>update documentation<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! The Resolver is responsible for performing recursive queries to lookup domain names.
//!
//! This is a 100% in process DNS resolver. It *does not* use the Host OS' resolver. If what is desired is to use the Host OS' resolver, generally in the system's libc, then the `std::net::ToSocketAddrs` variant over `&str` should be used.
//!
//! Unlike the `trust-dns` client, this tries to provide a simpler interface to perform DNS queries. For update options, i.e. Dynamic DNS, the `trust-dns` crate must be used directly. The Resolver library is capable of searching multiple domains (this can be disabled by using an FQDN during lookup), dual-stack IPv4/IPv6 lookups, performing chained CNAME lookups, and features connection metric tracking for attempting to pick the best upstream DNS resolver.
//!
//! There are two types for performing DNS queries, `Resolver` and `ResolverFuture`. `Resolver` is the easiest to work with, it is a wrapper around `ResolverFuture`. `ResolverFuture` is a `Tokio` based async resolver, and can be used inside any `Tokio` based system.
//!
//! This as best as possible attempts to abide by the the DNS RFCs, please file issues at https://github.com/bluejekyll/trust-dns .
//!
//! # Usage
//!
//! ## Declare dependency
//!
//! ```toml
//! [dependency]
//! trust-dns-resolver = "^0.6"
//! ```
//!
//! ## Extern the crate for usage in the library
//!
//! ```rust
//! extern crate trust_dns_resolver;
//! ```
//!
//! ## Using the Synchronous Resolver
//!
//! This uses the default configuration, which sets the [Google Public DNS](https://developers.google.com/speed/public-dns/) as the upstream resolvers. Please see their [privacy statement](https://developers.google.com/speed/public-dns/privacy) for important information about what they track, many ISP's track similar information in DNS.
//!
//! ```rust
//! use std::net::*;
//! use trust_dns_resolver::Resolver;
//! use trust_dns_resolver::config::*;
//!
//! // Construct a new Resolver with default configuration options
//! let resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();
//!
//! // Lookup the IP addresses associated with a name.
//! // The final dot forces this to be an FQDN, otherwise the search rules as specified
//! //  in `ResolverOpts` will take effect. FQDN's are generally cheaper queries.
//! let response = resolver.lookup_ip("www.example.com.").unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.iter().next().expect("no addresses returned!");
//! if address.is_ipv4() {
//!     assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! } else {
//!     assert_eq!(address, IpAddr::V6(Ipv6Addr::new(0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946)));
//! }
//! ```
//!
//! ## Using the host system config
//!
//! On Unix systems, the `/etc/resolv.conf` can be used for configuration. Not all options specified in the host systems `resolv.conf` are applicable or compatible with this software. In addition there may be additional options supported which the host system does not. Example:
//!
//! ```rust,no_run
//! # use std::net::*;
//! # use trust_dns_resolver::Resolver;
//! // Use the host OS'es `/etc/resolv.conf`
//! # #[cfg(unix)]
//! let resolver = Resolver::from_system_conf().unwrap();
//! # #[cfg(unix)]
//! let response = resolver.lookup_ip("www.example.com.").unwrap();
//! ```
//!
//! ## Using the Tokio/Async Resolver
//!
//! For more advanced asynchronous usage, the ResolverFuture is integrated with Tokio. In fact, the ResolverFuture is used by the synchronous Resolver for all lookups.
//!
//! ```rust
//! # extern crate futures;
//! # extern crate tokio_core;
//! # extern crate trust_dns_resolver;
//! # fn main() {
//! use std::net::*;
//! use tokio_core::reactor::Core;
//! use trust_dns_resolver::ResolverFuture;
//! use trust_dns_resolver::config::*;
//!
//! // We need a Tokio reactor::Core to run the resolver
//! //  this is responsible for running all Future tasks and registering interest in IO channels
//! let mut io_loop = Core::new().unwrap();
//!
//! // Construct a new Resolver with default configuration options
//! let resolver = ResolverFuture::new(ResolverConfig::default(), ResolverOpts::default(), &io_loop.handle());
//!
//! // Lookup the IP addresses associated with a name.
//! // This returns a future that will lookup the IP addresses, it must be run in the Core to
//! //  to get the actual result.
//! let lookup_future = resolver.lookup_ip("www.example.com.");
//!
//! // Run the lookup until it resolves or errors
//! let mut response = io_loop.run(lookup_future).unwrap();
//!
//! // There can be many addresses associated with the name,
//! //  this can return IPv4 and/or IPv6 addresses
//! let address = response.iter().next().expect("no addresses returned!");
//! if address.is_ipv4() {
//!     assert_eq!(address, IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)));
//! } else {
//!     assert_eq!(address, IpAddr::V6(Ipv6Addr::new(0x2606, 0x2800, 0x220, 0x1, 0x248, 0x1893, 0x25c8, 0x1946)));
//! }
//! # }
//! ```
//!
//! Generally after a lookup in an asynchornous context, there would probably be a connection made to a server, for example:
//!
//! ```c
//! let result = io_loop.run(lookup_future.and_then(|ips| {
//!                                  let ip = ips.next().unwrap();
//!                                  TcpStream::connect()
//!                              }).and_then(|conn| /* do something with the connection... */)
//!                          ).unwrap();
//! ```
//!
//! It's beyond the scope of these examples to show how to deal with connection failures and looping etc. But if you wanted to say try a different address from the result set after a connection failure, it will be necessary to create a type that implements the `Future` trait. Inside the `Future::poll` method would be the place to implement a loop over the different IP addresses.

#![deny(missing_docs)]

#[macro_use]
extern crate error_chain;
extern crate futures;
extern crate lalrpop_util;
#[macro_use]
extern crate log;
extern crate lru_cache;
extern crate tokio_core;
extern crate trust_dns_proto;
#[cfg(all(feature = "ipconfig", target_os = "windows", target_pointer_width = "64"))]
extern crate ipconfig;

pub mod config;
pub mod error;
pub mod lookup_ip;
pub mod lookup;
pub mod lookup_state;
#[doc(hidden)]
pub mod name_server_pool;
mod resolver;
pub mod system_conf;
mod resolver_future;
mod hosts;

pub use resolver::Resolver;
pub use resolver_future::ResolverFuture;
pub use hosts::Hosts;

/// returns a version as specified in Cargo.toml
pub fn version() -> &'static str {
    env!("CARGO_PKG_VERSION")
}
<|endoftext|>


<commit_before>use chrono::Utc;
use conduit_cookie::RequestSession;

use super::prelude::*;

use crate::middleware::log_request;
use crate::models::{ApiToken, User};
use crate::util::errors::{
    account_locked, forbidden, internal, AppError, AppResult, ChainError,
    InsecurelyGeneratedTokenRevoked,
};

#[derive(Debug)]
pub struct AuthenticatedUser {
    user: User,
    token_id: Option<i32>,
}

impl AuthenticatedUser {
    pub fn user_id(&self) -> i32 {
        self.user.id
    }

    pub fn api_token_id(&self) -> Option<i32> {
        self.token_id
    }

    pub fn user(self) -> User {
        self.user
    }
}

/// The Origin header (https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Origin)
/// is sent with CORS requests and POST requests, and indicates where the request comes from.
/// We don't want to accept authenticated requests that originated from other sites, so this
/// function returns an error if the Origin header doesn't match what we expect "this site" to
/// be: https://crates.io in production, or http://localhost:port/ in development.
fn verify_origin(req: &dyn RequestExt) -> AppResult<()> {
    let headers = req.headers();
    let allowed_origins = &req.app().config.allowed_origins;

    let bad_origin = headers
        .get_all(header::ORIGIN)
        .iter()
        .filter_map(|value| value.to_str().ok())
        .find(|value| !allowed_origins.iter().any(|it| it == value));

    if let Some(bad_origin) = bad_origin {
        let error_message = format!(
            "only same-origin requests can be authenticated. got {:?}",
            bad_origin
        );
        return Err(internal(&error_message))
            .chain_error(|| Box::new(forbidden()) as Box<dyn AppError>);
    }
    Ok(())
}

fn authenticate_user(req: &dyn RequestExt) -> AppResult<AuthenticatedUser> {
    let conn = req.db_conn()?;

    let session = req.session();
    let user_id_from_session = session.get("user_id").and_then(|s| s.parse::<i32>().ok());

    let (user_id, token_id) = if let Some(id) = user_id_from_session {
        (id, None)
    } else {
        // Otherwise, look for an `Authorization` header on the request
        let maybe_authorization: Option<String> = {
            req.headers()
                .get(header::AUTHORIZATION)
                .and_then(|h| h.to_str().ok())
                .map(|h| h.to_string())
        };
        if let Some(header_value) = maybe_authorization {
            let token = ApiToken::find_by_api_token(&conn, &header_value).map_err(|e| {
                if e.is::<InsecurelyGeneratedTokenRevoked>() {
                    e
                } else {
                    e.chain(internal("invalid token")).chain(forbidden())
                }
            })?;

            (token.user_id, Some(token.id))
        } else {
            // Unable to authenticate the user
            return Err(internal("no cookie session or auth header found")).chain_error(forbidden);
        }
    };

    let user = User::find(&conn, user_id)
        .chain_error(|| internal("user_id from cookie or token not found in database"))?;

    Ok(AuthenticatedUser { user, token_id })
}

impl<'a> UserAuthenticationExt for dyn RequestExt + 'a {
    /// Obtain `AuthenticatedUser` for the request or return an `Forbidden` error
    fn authenticate(&mut self) -> AppResult<AuthenticatedUser> {
        verify_origin(self)?;

        let authenticated_user = authenticate_user(self)?;

        if let Some(reason) = &authenticated_user.user.account_lock_reason {
            let still_locked = if let Some(until) = authenticated_user.user.account_lock_until {
                until > Utc::now().naive_utc()
            } else {
                true
            };
            if still_locked {
                return Err(account_locked(
                    &reason,
                    authenticated_user.user.account_lock_until,
                ));
            }
        }

        log_request::add_custom_metadata(self, "uid", authenticated_user.user_id());
        if let Some(id) = authenticated_user.api_token_id() {
            log_request::add_custom_metadata(self, "tokenid", id);
        }

        Ok(authenticated_user)
    }
}
<commit_msg>controllers::util: Simplify `maybe_authorization` variable assignment<commit_after>use chrono::Utc;
use conduit_cookie::RequestSession;

use super::prelude::*;

use crate::middleware::log_request;
use crate::models::{ApiToken, User};
use crate::util::errors::{
    account_locked, forbidden, internal, AppError, AppResult, ChainError,
    InsecurelyGeneratedTokenRevoked,
};

#[derive(Debug)]
pub struct AuthenticatedUser {
    user: User,
    token_id: Option<i32>,
}

impl AuthenticatedUser {
    pub fn user_id(&self) -> i32 {
        self.user.id
    }

    pub fn api_token_id(&self) -> Option<i32> {
        self.token_id
    }

    pub fn user(self) -> User {
        self.user
    }
}

/// The Origin header (https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Origin)
/// is sent with CORS requests and POST requests, and indicates where the request comes from.
/// We don't want to accept authenticated requests that originated from other sites, so this
/// function returns an error if the Origin header doesn't match what we expect "this site" to
/// be: https://crates.io in production, or http://localhost:port/ in development.
fn verify_origin(req: &dyn RequestExt) -> AppResult<()> {
    let headers = req.headers();
    let allowed_origins = &req.app().config.allowed_origins;

    let bad_origin = headers
        .get_all(header::ORIGIN)
        .iter()
        .filter_map(|value| value.to_str().ok())
        .find(|value| !allowed_origins.iter().any(|it| it == value));

    if let Some(bad_origin) = bad_origin {
        let error_message = format!(
            "only same-origin requests can be authenticated. got {:?}",
            bad_origin
        );
        return Err(internal(&error_message))
            .chain_error(|| Box::new(forbidden()) as Box<dyn AppError>);
    }
    Ok(())
}

fn authenticate_user(req: &dyn RequestExt) -> AppResult<AuthenticatedUser> {
    let conn = req.db_conn()?;

    let session = req.session();
    let user_id_from_session = session.get("user_id").and_then(|s| s.parse::<i32>().ok());

    let (user_id, token_id) = if let Some(id) = user_id_from_session {
        (id, None)
    } else {
        // Otherwise, look for an `Authorization` header on the request
        let maybe_authorization = req
            .headers()
            .get(header::AUTHORIZATION)
            .and_then(|h| h.to_str().ok());

        if let Some(header_value) = maybe_authorization {
            let token = ApiToken::find_by_api_token(&conn, header_value).map_err(|e| {
                if e.is::<InsecurelyGeneratedTokenRevoked>() {
                    e
                } else {
                    e.chain(internal("invalid token")).chain(forbidden())
                }
            })?;

            (token.user_id, Some(token.id))
        } else {
            // Unable to authenticate the user
            return Err(internal("no cookie session or auth header found")).chain_error(forbidden);
        }
    };

    let user = User::find(&conn, user_id)
        .chain_error(|| internal("user_id from cookie or token not found in database"))?;

    Ok(AuthenticatedUser { user, token_id })
}

impl<'a> UserAuthenticationExt for dyn RequestExt + 'a {
    /// Obtain `AuthenticatedUser` for the request or return an `Forbidden` error
    fn authenticate(&mut self) -> AppResult<AuthenticatedUser> {
        verify_origin(self)?;

        let authenticated_user = authenticate_user(self)?;

        if let Some(reason) = &authenticated_user.user.account_lock_reason {
            let still_locked = if let Some(until) = authenticated_user.user.account_lock_until {
                until > Utc::now().naive_utc()
            } else {
                true
            };
            if still_locked {
                return Err(account_locked(
                    &reason,
                    authenticated_user.user.account_lock_until,
                ));
            }
        }

        log_request::add_custom_metadata(self, "uid", authenticated_user.user_id());
        if let Some(id) = authenticated_user.api_token_id() {
            log_request::add_custom_metadata(self, "tokenid", id);
        }

        Ok(authenticated_user)
    }
}
<|endoftext|>


<commit_before>#[cfg(feature="no_std")]
use core::{cell, mem};
#[cfg(feature="no_std")]
#[cfg_attr(feature="no_std", allow(unused_imports))] // HACK: seems to be a bug in this lint!
use core::num::Float;
#[cfg(not(feature="no_std"))]
use std::{cell, mem};
#[cfg(feature="use_complex")]
use num_complex::Complex;

use AbsDiffEq;

/// Equality comparisons between two numbers using both the absolute difference and ULPs
/// (Units in Last Place) based comparisons.
pub trait UlpsEq: AbsDiffEq {
    /// The default ULPs to tolerate when testing values that are far-apart.
    ///
    /// This is used when no `max_ulps` value is supplied to the `ulps_eq` macro.
    fn default_max_ulps() -> u32;

    /// A test for equality that uses units in the last place (ULP) if the values are far apart.
    fn ulps_eq(&self, other: &Self, epsilon: Self::Epsilon, max_ulps: u32) -> bool;

    /// The inverse of `ApproxEq::ulps_eq`.
    fn ulps_ne(&self, other: &Self, epsilon: Self::Epsilon, max_ulps: u32) -> bool {
        !Self::ulps_eq(self, other, epsilon, max_ulps)
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Base implementations
///////////////////////////////////////////////////////////////////////////////////////////////////

// Implementation based on: [Comparing Floating Point Numbers, 2012 Edition]
// (https://randomascii.wordpress.com/2012/02/25/comparing-floating-point-numbers-2012-edition/)
macro_rules! impl_ulps_eq {
    ($T:ident, $U:ident) => {
        impl UlpsEq for $T {
            #[inline]
            fn default_max_ulps() -> u32 { 4 }

            #[inline]
            fn ulps_eq(&self, other: &$T, epsilon: $T, max_ulps: u32) -> bool {
                // For when the numbers are really close together
                if $T::abs_diff_eq(self, other, epsilon) {
                    return true;
                }

                // Trivial negative sign check
                if self.signum() != other.signum() {
                    return false;
                }

                // ULPS difference comparison
                let int_self: $U = unsafe { mem::transmute(*self) };
                let int_other: $U = unsafe { mem::transmute(*other) };

                $U::abs(int_self - int_other) < max_ulps as $U
            }
        }
    }
}

impl_ulps_eq!(f32, i32);
impl_ulps_eq!(f64, i64);

///////////////////////////////////////////////////////////////////////////////////////////////////
// Derived implementations
///////////////////////////////////////////////////////////////////////////////////////////////////

impl<'a, T: UlpsEq> UlpsEq for &'a T {
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &&'a T, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(*self, *other, epsilon, max_ulps)
    }
}

impl<'a, T: UlpsEq> UlpsEq for &'a mut T {
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &&'a mut T, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(*self, *other, epsilon, max_ulps)
    }
}

impl<T: UlpsEq + Copy> UlpsEq for cell::Cell<T> {
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &cell::Cell<T>, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(&self.get(), &other.get(), epsilon, max_ulps)
    }
}

impl<T: UlpsEq> UlpsEq for cell::RefCell<T> {
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &cell::RefCell<T>, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(&self.borrow(), &other.borrow(), epsilon, max_ulps)
    }
}

impl<T: UlpsEq> UlpsEq for [T]
    where T::Epsilon: Clone
{
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &[T], epsilon: T::Epsilon, max_ulps: u32) -> bool {
        self.len() == other.len() &&
        Iterator::zip(self.iter(), other).all(|(x, y)| {
                                                  T::ulps_eq(x,
                                                             y,
                                                             epsilon.clone(),
                                                             max_ulps.clone())
                                              })
    }
}

#[cfg(feature="use_complex")]
impl<T: UlpsEq> UlpsEq for Complex<T>
    where T::Epsilon: Clone
{
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &Complex<T>, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(&self.re, &other.re, epsilon.clone(), max_ulps) &&
        T::ulps_eq(&self.im, &other.im, epsilon.clone(), max_ulps)
    }
}
<commit_msg>Change ulps_eq() to use '<=' comparison<commit_after>#[cfg(feature="no_std")]
use core::{cell, mem};
#[cfg(feature="no_std")]
#[cfg_attr(feature="no_std", allow(unused_imports))] // HACK: seems to be a bug in this lint!
use core::num::Float;
#[cfg(not(feature="no_std"))]
use std::{cell, mem};
#[cfg(feature="use_complex")]
use num_complex::Complex;

use AbsDiffEq;

/// Equality comparisons between two numbers using both the absolute difference and ULPs
/// (Units in Last Place) based comparisons.
pub trait UlpsEq: AbsDiffEq {
    /// The default ULPs to tolerate when testing values that are far-apart.
    ///
    /// This is used when no `max_ulps` value is supplied to the `ulps_eq` macro.
    fn default_max_ulps() -> u32;

    /// A test for equality that uses units in the last place (ULP) if the values are far apart.
    fn ulps_eq(&self, other: &Self, epsilon: Self::Epsilon, max_ulps: u32) -> bool;

    /// The inverse of `ApproxEq::ulps_eq`.
    fn ulps_ne(&self, other: &Self, epsilon: Self::Epsilon, max_ulps: u32) -> bool {
        !Self::ulps_eq(self, other, epsilon, max_ulps)
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Base implementations
///////////////////////////////////////////////////////////////////////////////////////////////////

// Implementation based on: [Comparing Floating Point Numbers, 2012 Edition]
// (https://randomascii.wordpress.com/2012/02/25/comparing-floating-point-numbers-2012-edition/)
macro_rules! impl_ulps_eq {
    ($T:ident, $U:ident) => {
        impl UlpsEq for $T {
            #[inline]
            fn default_max_ulps() -> u32 { 4 }

            #[inline]
            fn ulps_eq(&self, other: &$T, epsilon: $T, max_ulps: u32) -> bool {
                // For when the numbers are really close together
                if $T::abs_diff_eq(self, other, epsilon) {
                    return true;
                }

                // Trivial negative sign check
                if self.signum() != other.signum() {
                    return false;
                }

                // ULPS difference comparison
                let int_self: $U = unsafe { mem::transmute(*self) };
                let int_other: $U = unsafe { mem::transmute(*other) };

                $U::abs(int_self - int_other) <= max_ulps as $U
            }
        }
    }
}

impl_ulps_eq!(f32, i32);
impl_ulps_eq!(f64, i64);

///////////////////////////////////////////////////////////////////////////////////////////////////
// Derived implementations
///////////////////////////////////////////////////////////////////////////////////////////////////

impl<'a, T: UlpsEq> UlpsEq for &'a T {
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &&'a T, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(*self, *other, epsilon, max_ulps)
    }
}

impl<'a, T: UlpsEq> UlpsEq for &'a mut T {
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &&'a mut T, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(*self, *other, epsilon, max_ulps)
    }
}

impl<T: UlpsEq + Copy> UlpsEq for cell::Cell<T> {
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &cell::Cell<T>, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(&self.get(), &other.get(), epsilon, max_ulps)
    }
}

impl<T: UlpsEq> UlpsEq for cell::RefCell<T> {
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &cell::RefCell<T>, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(&self.borrow(), &other.borrow(), epsilon, max_ulps)
    }
}

impl<T: UlpsEq> UlpsEq for [T]
    where T::Epsilon: Clone
{
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &[T], epsilon: T::Epsilon, max_ulps: u32) -> bool {
        self.len() == other.len() &&
        Iterator::zip(self.iter(), other).all(|(x, y)| {
                                                  T::ulps_eq(x,
                                                             y,
                                                             epsilon.clone(),
                                                             max_ulps.clone())
                                              })
    }
}

#[cfg(feature="use_complex")]
impl<T: UlpsEq> UlpsEq for Complex<T>
    where T::Epsilon: Clone
{
    #[inline]
    fn default_max_ulps() -> u32 {
        T::default_max_ulps()
    }

    #[inline]
    fn ulps_eq(&self, other: &Complex<T>, epsilon: T::Epsilon, max_ulps: u32) -> bool {
        T::ulps_eq(&self.re, &other.re, epsilon.clone(), max_ulps) &&
        T::ulps_eq(&self.im, &other.im, epsilon.clone(), max_ulps)
    }
}
<|endoftext|>