<commit_before>use std::cmp::Ordering;
use std::collections::HashSet;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};

pub type Tag = u32;
pub type Offset = u32;
pub type Generation = u32;

/// Our default value for uninitialized or unconnected components in the mesh.
pub const INVALID_ELEMENT_OFFSET: Offset = 0;
pub const INVALID_ELEMENT_GENERATION: Generation = 0;

/// Type-safe index into kernel storage.
#[derive(Default, Debug, Eq)]
pub struct Handle<T> {
    pub offset: Offset,
    pub generation: Generation,
    _marker: PhantomData<T>,
}

impl<T> Copy for Handle<T> {}
impl<T> Clone for Handle<T> {
    fn clone(&self) -> Self {
        Self {
            offset: self.offset,
            generation: self.generation,
            _marker: self._marker,
        }
    }
}

impl<T> Hash for Handle<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.offset.hash(state);
        self.generation.hash(state);
    }
}

impl<T> Handle<T> {
    pub fn new(offset: Offset, generation: Generation) -> Handle<T> {
        Handle {
            offset,
            generation,
            _marker: PhantomData::default(),
        }
    }

    pub fn is_valid(&self) -> bool {
        self.offset != INVALID_ELEMENT_OFFSET
    }
}

impl<T> PartialOrd for Handle<T> {
    fn partial_cmp(&self, other: &Handle<T>) -> Option<Ordering> {
        // Only the offset should matter when it comes to ordering
        self.offset.partial_cmp(&other.offset)
    }
}

impl<T> PartialEq for Handle<T> {
    fn eq(&self, other: &Handle<T>) -> bool {
        self.offset.eq(&other.offset) && self.generation.eq(&other.generation)
    }
}

/// A pretty simple wrapper over a pair of 'Vec's.
pub struct ElementBuffer<D: Default> {
    buffer: Vec<D>,
    generations: Vec<Generation>,
    // Why not put the index? Because the generation of an index could give us
    // false negatives if we're not careful ... I'm still considering this.
    free_cells: HashSet<Offset>,
    //tags: Vec<Tag>, // TODO: use a Set instead. This isn't a persistent array of attributes.
}

impl<D: Default> Default for ElementBuffer<D> {
    fn default() -> Self {
        ElementBuffer {
            buffer: vec![Default::default()],
            generations: vec![Default::default()],
            free_cells: HashSet::new(),
            //tags: Vec::new(),
        }
    }
}

impl<D: Default> fmt::Debug for ElementBuffer<D> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ElementBuffer<> {{ {} items }}", self.len())
    }
}

impl<D: Default> ElementBuffer<D> {
    #[inline(always)]
    fn is_active_cell(&self, offset: Offset) -> bool {
        !self.free_cells.contains(&offset)
    }

    /// Returns the number of currently active cells.
    /// The actual number of items allocated by the buffer might
    /// be different.
    pub fn len(&self) -> usize {
        (self.buffer.len() - 1) - self.free_cells.len()
    }

    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn has_inactive_cells(&self) -> bool {
        !self.free_cells.is_empty()
    }

    pub fn iter(&self) -> impl Iterator<Item = (Handle<D>, &D)> {
        self.buffer
            .iter()
            .enumerate()
            .skip(1)
            .zip(self.generations.iter().skip(1))
            .filter(|((offset, _), _)| self.is_active_cell(*offset as Offset))
            .map(|((offset, element), generation)| {
                (Handle::new(offset as Offset, *generation), element)
            })
    }

    pub fn iter_mut(&mut self) -> impl Iterator<Item = (Handle<D>, &mut D)> {
        self.buffer
            .iter_mut()
            .enumerate()
            .skip(1)
            .zip(self.generations.iter().skip(1))
            .filter(|((offset, _), _)| {
                let offset = *offset as Offset;
                !self.free_cells.contains(&offset)
            })
            .map(|((offset, element), generation)| {
                (Handle::new(offset as Offset, *generation), element)
            })
    }

    pub fn get(&self, handle: Handle<D>) -> Option<&D> {
        if !self.is_active_cell(handle.offset) {
            return None;
        }

        let generation = self.generations[handle.offset as usize];
        if generation != handle.generation {
            return None;
        }

        self.buffer.get(handle.offset as usize)
    }

    pub fn get_mut(&mut self, handle: Handle<D>) -> Option<&mut D> {
        if !self.is_active_cell(handle.offset) {
            return None;
        }

        let generation = self.generations[handle.offset as usize];
        if generation != handle.generation {
            return None;
        }

        self.buffer.get_mut(handle.offset as usize)
    }

    pub fn push(&mut self, element: D) -> Handle<D> {
        if let Some(offset) = self.free_cells.iter().next().cloned() {
            self.free_cells.remove(&offset);
            // In this situation we just re-use an existing cell
            self.buffer[offset as usize] = element;
            Handle::new(offset, self.generations[offset as usize])
        } else {
            // Here we push the element on to the back
            let offset = self.buffer.len() as Offset;
            self.buffer.push(element);
            self.generations.push(1);
            Handle::new(offset, 1)
        }
    }

    pub fn remove(&mut self, handle: Handle<D>) {
        self.free_cells.insert(handle.offset);
        self.generations[handle.offset as usize] += 1;
    }

    // fn next_swap_pair(&self) -> Option<(Offset, Offset)> {
    //     let inactive_offset = self.enumerate().find(|e| !e.1.is_active()).map(|e| e.0);
    //     let active_offset = self
    //         .enumerate()
    //         .rev()
    //         .find(|e| e.1.is_active())
    //         .map(|e| e.0);
    //     if let (Some(inactive_offset), Some(active_offset)) = (inactive_offset, active_offset) {
    //         if active_offset < inactive_offset {
    //             log::debug!("Buffer appears to be successfully sorted!");
    //             // by the time this is true we should have sorted/swapped
    //             // all elements so that the inactive inactive elements
    //             // make up the tail of the buffer.
    //             None
    //         } else {
    //             Some((inactive_offset as u32, active_offset as u32))
    //         }
    //     } else {
    //         log::debug!("No more swap pairs.");
    //         None
    //     }
    // }
}

impl<D: Default> Index<Handle<D>> for ElementBuffer<D> {
    type Output = D;

    fn index(&self, handle: Handle<D>) -> &Self::Output {
        self.get(handle)
            .expect("Unable to retrieve item for provided handle.")
    }
}

impl<D: Default> IndexMut<Handle<D>> for ElementBuffer<D> {
    fn index_mut(&mut self, handle: Handle<D>) -> &mut Self::Output {
        self.get_mut(handle).expect("Unable to retrieve item.")
    }
}

///////////////////////////////////////////////////////////////////////////////

pub mod prelude {
    pub use super::{Generation, Handle, Tag};
}

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

    #[derive(Default)]
    struct TestElement {
        foo: u32,
    }

    type TestHandle = Handle<TestElement>;
    type TestBuffer = ElementBuffer<TestElement>;

    #[test]
    fn default_index_is_invalid() {
        let index = TestHandle::default();
        assert!(!index.is_valid());
    }

    #[test]
    fn default_element_buffer_properties() {
        let buffer = TestBuffer::default();
        assert!(!buffer.has_inactive_cells());
        assert_eq!(buffer.len(), 0);
        assert!(buffer.is_empty());
    }

    #[test]
    fn freecells_set_test() {
        let mut foo: HashSet<u32> = HashSet::new();
        foo.insert(1);
        foo.insert(10);
        foo.insert(256);

        assert_eq!(foo.len(), 3);
        assert!(foo.iter().next().is_some());
        assert_eq!(foo.len(), 3);

        let offset = foo
            .iter()
            .next()
            .cloned()
            .expect("Failed to get first value from set.");
        assert!(foo.contains(&offset));
        foo.remove(&offset);
        assert!(!foo.contains(&offset));

        assert_eq!(foo.len(), 2);
    }

    #[test]
    fn push_elements() {
        let mut buffer = TestBuffer::default();

        assert!(buffer.is_empty());

        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });

        assert!(!buffer.is_empty());
        assert_eq!(buffer.len(), 2);

        assert_eq!(i0.offset, 1);
        assert_eq!(i0.generation, 1);

        assert_eq!(i1.offset, 2);
        assert_eq!(i1.generation, 1);
    }

    #[test]
    fn iter_elements() {
        let mut buffer = TestBuffer::default();
        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });
        let i2 = buffer.push(TestElement { foo: 2 });

        assert_eq!(buffer.iter().count(), 3);

        {
            let offsets: Vec<Offset> = buffer.iter().map(|(index, _)| index.offset).collect();
            assert_eq!(offsets.len(), 3);
            assert_eq!(offsets[0], i0.offset);
            assert_eq!(offsets[1], i1.offset);
            assert_eq!(offsets[2], i2.offset);
        }

        {
            let mut it = buffer.iter();

            let (_, element) = it.next().expect("Unexpected end of iterator");
            assert_eq!(element.foo, 0);

            let (_, element) = it.next().expect("Unexpected end of iterator");
            assert_eq!(element.foo, 1);

            let (_, element) = it.next().expect("Unexpected end of iterator");
            assert_eq!(element.foo, 2);

            assert!(it.next().is_none());
        }

        {
            for (_, element) in buffer.iter_mut() {
                element.foo += 1;
            }

            let foos: Vec<u32> = buffer.iter().map(|(_, e)| e.foo).collect();
            assert_eq!(foos[0], 1);
            assert_eq!(foos[1], 2);
            assert_eq!(foos[2], 3);
        }
    }

    #[test]
    fn indexing() {
        let mut buffer = TestBuffer::default();
        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });
        let i2 = buffer.push(TestElement { foo: 2 });

        assert_eq!(buffer[i0].foo, 0);
        assert_eq!(buffer[i1].foo, 1);
        assert_eq!(buffer[i2].foo, 2);

        buffer[i2].foo = 3;

        let foo_accum: u32 = buffer.iter().map(|(_, e)| e.foo).sum();
        assert_eq!(foo_accum, 4);
    }

    #[test]
    fn remove_elements() {
        let mut buffer = TestBuffer::default();
        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });
        let i2 = buffer.push(TestElement { foo: 2 });
        let i3 = buffer.push(TestElement { foo: 3 });
        let i4 = buffer.push(TestElement { foo: 4 });

        assert_eq!(buffer.len(), 5);

        buffer.remove(i2);
        buffer.remove(i3);

        assert_eq!(buffer.len(), 3);

        let foos: Vec<u32> = buffer.iter().map(|(_, e)| e.foo).collect();
        assert_eq!(foos[0], 0);
        assert_eq!(foos[1], 1);
        assert_eq!(foos[2], 4);

        let offsets: Vec<Offset> = buffer.iter().map(|(h, _)| h.offset).collect();
        assert_eq!(offsets[0], i0.offset);
        assert_eq!(offsets[1], i1.offset);
        assert_eq!(offsets[2], i4.offset);

        assert!(buffer.get(i2).is_none());
        assert!(buffer.get(i3).is_none());
    }

    #[test]
    fn insert_after_remove() {
        let mut buffer = TestBuffer::default();
        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });
        let i2 = buffer.push(TestElement { foo: 2 });
        let i3 = buffer.push(TestElement { foo: 3 });
        let i4 = buffer.push(TestElement { foo: 4 });

        assert_eq!(i0.generation, 1);
        assert_eq!(i1.generation, 1);
        assert_eq!(i2.generation, 1);
        assert_eq!(i3.generation, 1);
        assert_eq!(i4.generation, 1);

        assert_eq!(buffer.len(), 5);

        buffer.remove(i2);
        buffer.remove(i3);

        assert_eq!(buffer.len(), 3);

        let i5 = buffer.push(TestElement { foo: 5 });
        let i6 = buffer.push(TestElement { foo: 6 });

        assert_eq!(i5.generation, 2);
        assert_eq!(i6.generation, 2);

        assert_eq!(buffer.len(), 5);

        assert!(buffer.get(i2).is_none());
        assert!(buffer.get(i3).is_none());

        assert!(buffer.get(i5).is_some());
        assert!(buffer.get(i6).is_some());

        assert_eq!(buffer[i5].foo, 5);
        assert_eq!(buffer[i6].foo, 6);
    }
}
<commit_msg>Adding with_capacity initializer<commit_after>use std::cmp::Ordering;
use std::collections::HashSet;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};

pub type Tag = u32;
pub type Offset = u32;
pub type Generation = u32;

/// Our default value for uninitialized or unconnected components in the mesh.
pub const INVALID_ELEMENT_OFFSET: Offset = 0;
pub const INVALID_ELEMENT_GENERATION: Generation = 0;

/// Type-safe index into kernel storage.
#[derive(Default, Debug, Eq)]
pub struct Handle<T> {
    pub offset: Offset,
    pub generation: Generation,
    _marker: PhantomData<T>,
}

impl<T> Copy for Handle<T> {}
impl<T> Clone for Handle<T> {
    fn clone(&self) -> Self {
        Self {
            offset: self.offset,
            generation: self.generation,
            _marker: self._marker,
        }
    }
}

impl<T> Hash for Handle<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.offset.hash(state);
        self.generation.hash(state);
    }
}

impl<T> Handle<T> {
    pub fn new(offset: Offset, generation: Generation) -> Handle<T> {
        Handle {
            offset,
            generation,
            _marker: PhantomData::default(),
        }
    }

    pub fn is_valid(&self) -> bool {
        self.offset != INVALID_ELEMENT_OFFSET
    }
}

impl<T> PartialOrd for Handle<T> {
    fn partial_cmp(&self, other: &Handle<T>) -> Option<Ordering> {
        // Only the offset should matter when it comes to ordering
        self.offset.partial_cmp(&other.offset)
    }
}

impl<T> PartialEq for Handle<T> {
    fn eq(&self, other: &Handle<T>) -> bool {
        self.offset.eq(&other.offset) && self.generation.eq(&other.generation)
    }
}

/// A pretty simple wrapper over a pair of 'Vec's.
pub struct ElementBuffer<D: Default> {
    buffer: Vec<D>,
    generations: Vec<Generation>,
    // Why not put the index? Because the generation of an index could give us
    // false negatives if we're not careful ... I'm still considering this.
    free_cells: HashSet<Offset>,
    //tags: Vec<Tag>, // TODO: use a Set instead. This isn't a persistent array of attributes.
}

impl<D: Default> Default for ElementBuffer<D> {
    fn default() -> Self {
        ElementBuffer {
            buffer: vec![Default::default()],
            generations: vec![Default::default()],
            free_cells: HashSet::new(),
            //tags: Vec::new(),
        }
    }
}

impl<D: Default> fmt::Debug for ElementBuffer<D> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ElementBuffer<> {{ {} items }}", self.len())
    }
}

impl<D: Default> ElementBuffer<D> {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_capacity(capacity: usize) -> Self {
        let mut out = Self {
            buffer: Vec::with_capacity(capacity + 1),
            generations: Vec::with_capacity(capacity + 1),
            free_cells: HashSet::new(),
        };
        out.buffer.push(Default::default());
        out.generations.push(Default::default());
        out
    }

    #[inline(always)]
    fn is_active_cell(&self, offset: Offset) -> bool {
        !self.free_cells.contains(&offset)
    }

    /// Returns the number of currently active cells.
    /// The actual number of items allocated by the buffer might
    /// be different.
    pub fn len(&self) -> usize {
        (self.buffer.len() - 1) - self.free_cells.len()
    }

    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn has_inactive_cells(&self) -> bool {
        !self.free_cells.is_empty()
    }

    pub fn iter(&self) -> impl Iterator<Item = (Handle<D>, &D)> {
        self.buffer
            .iter()
            .enumerate()
            .skip(1)
            .zip(self.generations.iter().skip(1))
            .filter(|((offset, _), _)| self.is_active_cell(*offset as Offset))
            .map(|((offset, element), generation)| {
                (Handle::new(offset as Offset, *generation), element)
            })
    }

    pub fn iter_mut(&mut self) -> impl Iterator<Item = (Handle<D>, &mut D)> {
        self.buffer
            .iter_mut()
            .enumerate()
            .skip(1)
            .zip(self.generations.iter().skip(1))
            .filter(|((offset, _), _)| {
                let offset = *offset as Offset;
                !self.free_cells.contains(&offset)
            })
            .map(|((offset, element), generation)| {
                (Handle::new(offset as Offset, *generation), element)
            })
    }

    pub fn get(&self, handle: Handle<D>) -> Option<&D> {
        if !self.is_active_cell(handle.offset) {
            return None;
        }

        let generation = self.generations[handle.offset as usize];
        if generation != handle.generation {
            return None;
        }

        self.buffer.get(handle.offset as usize)
    }

    pub fn get_mut(&mut self, handle: Handle<D>) -> Option<&mut D> {
        if !self.is_active_cell(handle.offset) {
            return None;
        }

        let generation = self.generations[handle.offset as usize];
        if generation != handle.generation {
            return None;
        }

        self.buffer.get_mut(handle.offset as usize)
    }

    pub fn push(&mut self, element: D) -> Handle<D> {
        if let Some(offset) = self.free_cells.iter().next().cloned() {
            self.free_cells.remove(&offset);
            // In this situation we just re-use an existing cell
            self.buffer[offset as usize] = element;
            Handle::new(offset, self.generations[offset as usize])
        } else {
            // Here we push the element on to the back
            let offset = self.buffer.len() as Offset;
            self.buffer.push(element);
            self.generations.push(1);
            Handle::new(offset, 1)
        }
    }

    pub fn remove(&mut self, handle: Handle<D>) {
        self.free_cells.insert(handle.offset);
        self.generations[handle.offset as usize] += 1;
    }

    // fn next_swap_pair(&self) -> Option<(Offset, Offset)> {
    //     let inactive_offset = self.enumerate().find(|e| !e.1.is_active()).map(|e| e.0);
    //     let active_offset = self
    //         .enumerate()
    //         .rev()
    //         .find(|e| e.1.is_active())
    //         .map(|e| e.0);
    //     if let (Some(inactive_offset), Some(active_offset)) = (inactive_offset, active_offset) {
    //         if active_offset < inactive_offset {
    //             log::debug!("Buffer appears to be successfully sorted!");
    //             // by the time this is true we should have sorted/swapped
    //             // all elements so that the inactive inactive elements
    //             // make up the tail of the buffer.
    //             None
    //         } else {
    //             Some((inactive_offset as u32, active_offset as u32))
    //         }
    //     } else {
    //         log::debug!("No more swap pairs.");
    //         None
    //     }
    // }
}

impl<D: Default> Index<Handle<D>> for ElementBuffer<D> {
    type Output = D;

    fn index(&self, handle: Handle<D>) -> &Self::Output {
        self.get(handle)
            .expect("Unable to retrieve item for provided handle.")
    }
}

impl<D: Default> IndexMut<Handle<D>> for ElementBuffer<D> {
    fn index_mut(&mut self, handle: Handle<D>) -> &mut Self::Output {
        self.get_mut(handle).expect("Unable to retrieve item.")
    }
}

///////////////////////////////////////////////////////////////////////////////

pub mod prelude {
    pub use super::{ElementBuffer, Generation, Handle, Offset, Tag};
}

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

    #[derive(Default)]
    struct TestElement {
        foo: u32,
    }

    type TestHandle = Handle<TestElement>;
    type TestBuffer = ElementBuffer<TestElement>;

    #[test]
    fn default_index_is_invalid() {
        let index = TestHandle::default();
        assert!(!index.is_valid());
    }

    #[test]
    fn default_element_buffer_properties() {
        let buffer = TestBuffer::default();
        assert!(!buffer.has_inactive_cells());
        assert_eq!(buffer.len(), 0);
        assert!(buffer.is_empty());
    }

    #[test]
    fn freecells_set_test() {
        let mut foo: HashSet<u32> = HashSet::new();
        foo.insert(1);
        foo.insert(10);
        foo.insert(256);

        assert_eq!(foo.len(), 3);
        assert!(foo.iter().next().is_some());
        assert_eq!(foo.len(), 3);

        let offset = foo
            .iter()
            .next()
            .cloned()
            .expect("Failed to get first value from set.");
        assert!(foo.contains(&offset));
        foo.remove(&offset);
        assert!(!foo.contains(&offset));

        assert_eq!(foo.len(), 2);
    }

    #[test]
    fn push_elements() {
        let mut buffer = TestBuffer::default();

        assert!(buffer.is_empty());

        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });

        assert!(!buffer.is_empty());
        assert_eq!(buffer.len(), 2);

        assert_eq!(i0.offset, 1);
        assert_eq!(i0.generation, 1);

        assert_eq!(i1.offset, 2);
        assert_eq!(i1.generation, 1);
    }

    #[test]
    fn iter_elements() {
        let mut buffer = TestBuffer::default();
        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });
        let i2 = buffer.push(TestElement { foo: 2 });

        assert_eq!(buffer.iter().count(), 3);

        {
            let offsets: Vec<Offset> = buffer.iter().map(|(index, _)| index.offset).collect();
            assert_eq!(offsets.len(), 3);
            assert_eq!(offsets[0], i0.offset);
            assert_eq!(offsets[1], i1.offset);
            assert_eq!(offsets[2], i2.offset);
        }

        {
            let mut it = buffer.iter();

            let (_, element) = it.next().expect("Unexpected end of iterator");
            assert_eq!(element.foo, 0);

            let (_, element) = it.next().expect("Unexpected end of iterator");
            assert_eq!(element.foo, 1);

            let (_, element) = it.next().expect("Unexpected end of iterator");
            assert_eq!(element.foo, 2);

            assert!(it.next().is_none());
        }

        {
            for (_, element) in buffer.iter_mut() {
                element.foo += 1;
            }

            let foos: Vec<u32> = buffer.iter().map(|(_, e)| e.foo).collect();
            assert_eq!(foos[0], 1);
            assert_eq!(foos[1], 2);
            assert_eq!(foos[2], 3);
        }
    }

    #[test]
    fn indexing() {
        let mut buffer = TestBuffer::default();
        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });
        let i2 = buffer.push(TestElement { foo: 2 });

        assert_eq!(buffer[i0].foo, 0);
        assert_eq!(buffer[i1].foo, 1);
        assert_eq!(buffer[i2].foo, 2);

        buffer[i2].foo = 3;

        let foo_accum: u32 = buffer.iter().map(|(_, e)| e.foo).sum();
        assert_eq!(foo_accum, 4);
    }

    #[test]
    fn remove_elements() {
        let mut buffer = TestBuffer::default();
        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });
        let i2 = buffer.push(TestElement { foo: 2 });
        let i3 = buffer.push(TestElement { foo: 3 });
        let i4 = buffer.push(TestElement { foo: 4 });

        assert_eq!(buffer.len(), 5);

        buffer.remove(i2);
        buffer.remove(i3);

        assert_eq!(buffer.len(), 3);

        let foos: Vec<u32> = buffer.iter().map(|(_, e)| e.foo).collect();
        assert_eq!(foos[0], 0);
        assert_eq!(foos[1], 1);
        assert_eq!(foos[2], 4);

        let offsets: Vec<Offset> = buffer.iter().map(|(h, _)| h.offset).collect();
        assert_eq!(offsets[0], i0.offset);
        assert_eq!(offsets[1], i1.offset);
        assert_eq!(offsets[2], i4.offset);

        assert!(buffer.get(i2).is_none());
        assert!(buffer.get(i3).is_none());
    }

    #[test]
    fn insert_after_remove() {
        let mut buffer = TestBuffer::default();
        let i0 = buffer.push(TestElement { foo: 0 });
        let i1 = buffer.push(TestElement { foo: 1 });
        let i2 = buffer.push(TestElement { foo: 2 });
        let i3 = buffer.push(TestElement { foo: 3 });
        let i4 = buffer.push(TestElement { foo: 4 });

        assert_eq!(i0.generation, 1);
        assert_eq!(i1.generation, 1);
        assert_eq!(i2.generation, 1);
        assert_eq!(i3.generation, 1);
        assert_eq!(i4.generation, 1);

        assert_eq!(buffer.len(), 5);

        buffer.remove(i2);
        buffer.remove(i3);

        assert_eq!(buffer.len(), 3);

        let i5 = buffer.push(TestElement { foo: 5 });
        let i6 = buffer.push(TestElement { foo: 6 });

        assert_eq!(i5.generation, 2);
        assert_eq!(i6.generation, 2);

        assert_eq!(buffer.len(), 5);

        assert!(buffer.get(i2).is_none());
        assert!(buffer.get(i3).is_none());

        assert!(buffer.get(i5).is_some());
        assert!(buffer.get(i6).is_some());

        assert_eq!(buffer[i5].foo, 5);
        assert_eq!(buffer[i6].foo, 6);
    }
}
<|endoftext|>


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

use std::os::unix::io::RawFd;
use handle_table::{HandleTable, Handle};
use sys::unix::FdObserver;
use nix::sys::epoll;

pub struct Reactor {
    pub ep: RawFd,
    pub observers: HandleTable<FdObserver>,
    events: Vec<epoll::EpollEvent>
}

impl Reactor {
    pub fn new() -> Result<Reactor, ::std::io::Error> {
        Ok(Reactor {
            ep: try_syscall!(epoll::epoll_create()),
            observers: HandleTable::new(),
            events: Vec::with_capacity(1024),
        })
    }

    pub fn run_once(&mut self, maybe_timeout: Option<::time::Duration>)
                    -> Result<(), ::std::io::Error>
    {
        let timeout = match maybe_timeout {
            Some(t) if t >= ::time::Duration::zero() => t.num_milliseconds() as isize,
            _ => -1,
        };

        let events = unsafe {
            let ptr = (&mut self.events[..]).as_mut_ptr();
            ::std::slice::from_raw_parts_mut(ptr, self.events.capacity())
        };

        let n = try_syscall!(epoll::epoll_wait(self.ep, events, timeout));

        unsafe { self.events.set_len(n); }

        for event in &self.events {
            let handle = Handle { val: event.data as usize };

            if event.events.contains(epoll::EPOLLIN) || event.events.contains(epoll::EPOLLHUP) ||
               event.events.contains(epoll::EPOLLERR)
            {
                match self.observers[handle].read_fulfiller.take() {
                    None => (),
                    Some(fulfiller) => {
                        fulfiller.fulfill(());
                    }
                }
            }

            if event.events.contains(epoll::EPOLLOUT) || event.events.contains(epoll::EPOLLHUP) ||
               event.events.contains(epoll::EPOLLERR)
            {
                match self.observers[handle].write_fulfiller.take() {
                    None => (),
                    Some(fulfiller) => {
                        fulfiller.fulfill(());
                    }
                }
            }
        }
        Ok(())
    }

    pub fn new_observer(&mut self, fd: RawFd) -> Result<Handle, ::std::io::Error> {
        let observer = FdObserver { read_fulfiller: None, write_fulfiller: None };
        let handle = self.observers.push(observer);

        let info = epoll::EpollEvent {
            events: epoll::EPOLLIN | epoll::EPOLLOUT,
            data: handle.val as u64
        };

        try_syscall!(epoll::epoll_ctl(self.ep, epoll::EpollOp::EpollCtlAdd, fd, &info));

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

use std::os::unix::io::RawFd;
use handle_table::{HandleTable, Handle};
use sys::unix::FdObserver;
use nix::sys::epoll;

pub struct Reactor {
    pub ep: RawFd,
    pub observers: HandleTable<FdObserver>,
    events: Vec<epoll::EpollEvent>
}

impl Reactor {
    pub fn new() -> Result<Reactor, ::std::io::Error> {
        Ok(Reactor {
            ep: try_syscall!(epoll::epoll_create()),
            observers: HandleTable::new(),
            events: Vec::with_capacity(1024),
        })
    }

    pub fn run_once(&mut self, maybe_timeout: Option<::time::Duration>)
                    -> Result<(), ::std::io::Error>
    {
        let timeout = match maybe_timeout {
            Some(t) if t >= ::time::Duration::zero() => t.num_milliseconds() as isize,
            _ => -1,
        };

        let events = unsafe {
            let ptr = (&mut self.events[..]).as_mut_ptr();
            ::std::slice::from_raw_parts_mut(ptr, self.events.capacity())
        };

        let n = try_syscall!(epoll::epoll_wait(self.ep, events, timeout));

        unsafe { self.events.set_len(n); }

        for event in &self.events {
            let handle = Handle { val: event.data as usize };

            if event.events.contains(epoll::EPOLLIN) || event.events.contains(epoll::EPOLLHUP) ||
               event.events.contains(epoll::EPOLLERR)
            {
                match self.observers[handle].read_fulfiller.take() {
                    None => (),
                    Some(fulfiller) => {
                        fulfiller.fulfill(());
                    }
                }
            }

            if event.events.contains(epoll::EPOLLOUT) || event.events.contains(epoll::EPOLLHUP) ||
               event.events.contains(epoll::EPOLLERR)
            {
                match self.observers[handle].write_fulfiller.take() {
                    None => (),
                    Some(fulfiller) => {
                        fulfiller.fulfill(());
                    }
                }
            }
        }
        Ok(())
    }

    pub fn new_observer(&mut self, fd: RawFd) -> Result<Handle, ::std::io::Error> {
        let observer = FdObserver { read_fulfiller: None, write_fulfiller: None };
        let handle = self.observers.push(observer);

        let info = epoll::EpollEvent {
            events: epoll::EPOLLIN | epoll::EPOLLOUT | epoll::EPOLLET,
            data: handle.val as u64
        };

        try_syscall!(epoll::epoll_ctl(self.ep, epoll::EpollOp::EpollCtlAdd, fd, &info));

        Ok(handle)
    }
}
<|endoftext|>


<commit_before>extern crate security_framework;
extern crate tempdir;

use self::security_framework::base;
use self::security_framework::certificate::SecCertificate;
use self::security_framework::identity::SecIdentity;
use self::security_framework::import_export::Pkcs12ImportOptions;
use self::security_framework::random::SecRandom;
use self::security_framework::secure_transport::{self, SslContext, ProtocolSide, ConnectionType};
use self::security_framework::os::macos::keychain;
use self::security_framework::os::macos::import_export::Pkcs12ImportOptionsExt;
use self::tempdir::TempDir;
use std::fmt;
use std::io;
use std::error;

pub struct Error(base::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<base::Error> for Error {
    fn from(error: base::Error) -> Error {
        Error(error)
    }
}

pub struct Certificate(SecCertificate);

pub struct Identity(SecIdentity);

pub struct Pkcs12 {
    pub identity: Identity,
    pub chain: Vec<Certificate>,
}

impl Pkcs12 {
    pub fn parse(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let dir = TempDir::new("native_tls").unwrap(); //fixme
        let keychain = try!(keychain::CreateOptions::new()
            .password(pass) // FIXME maybe generate a secure random password here?
            .create(dir.path().join("keychain")));

        let mut import = try!(Pkcs12ImportOptions::new()
            .passphrase(pass)
            .keychain(keychain)
            .import(buf));
        let import = import.pop().unwrap();

        // The identity's cert shows up in the chain, so filter it out to avoid sending twice
        // FIXME should probably use CFEquals here
        let identity_cert = try!(import.identity.certificate()).to_der();

        Ok(Pkcs12 {
            identity: Identity(import.identity),
            chain: import.cert_chain
                .into_iter()
                .filter(|c| c.to_der() != identity_cert)
                .map(Certificate)
                .collect(),
        })
    }
}

pub enum HandshakeError<S> {
    Interrupted(MidHandshakeTlsStream<S>),
    Failure(Error),
}

impl<S> From<secure_transport::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<base::Error> for HandshakeError<S> {
    fn from(e: base::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub struct MidHandshakeTlsStream<S>(secure_transport::MidHandshakeSslStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
    where S: fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S>
    where S: io::Read + io::Write
{
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct ClientBuilder(());

impl ClientBuilder {
    pub fn new() -> Result<ClientBuilder, Error> {
        Ok(ClientBuilder(()))
    }

    pub fn handshake<S>(&mut self,
                        domain: &str,
                        stream: S)
                        -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Client, ConnectionType::Stream));
        try!(ctx.set_peer_domain_name(domain));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct ServerBuilder {
    identity: SecIdentity,
    chain: Vec<SecCertificate>,
}

impl ServerBuilder {
    pub fn new<I>(identity: Identity, chain: I) -> Result<ServerBuilder, Error>
        where I: IntoIterator<Item = Certificate>
    {
        Ok(ServerBuilder {
            identity: identity.0,
            chain: chain.into_iter().map(|c| c.0).collect(),
        })
    }

    pub fn handshake<S>(&mut self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Server, ConnectionType::Stream));
        try!(ctx.set_certificate(&self.identity, &self.chain));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(secure_transport::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(try!(self.0.context().buffered_read_size()))
    }
}

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

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

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

/// Security Framework-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the Security Framework `SslStream`.
    fn raw_stream(&self) -> &secure_transport::SslStream<S>;

    /// Returns a mutable reference to the Security Framework `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &secure_transport::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S> {
        &mut (self.0).0
    }
}
<commit_msg>One more fixme<commit_after>extern crate security_framework;
extern crate tempdir;

use self::security_framework::base;
use self::security_framework::certificate::SecCertificate;
use self::security_framework::identity::SecIdentity;
use self::security_framework::import_export::Pkcs12ImportOptions;
use self::security_framework::random::SecRandom;
use self::security_framework::secure_transport::{self, SslContext, ProtocolSide, ConnectionType};
use self::security_framework::os::macos::keychain;
use self::security_framework::os::macos::import_export::Pkcs12ImportOptionsExt;
use self::tempdir::TempDir;
use std::fmt;
use std::io;
use std::error;

pub struct Error(base::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<base::Error> for Error {
    fn from(error: base::Error) -> Error {
        Error(error)
    }
}

pub struct Certificate(SecCertificate);

pub struct Identity(SecIdentity);

pub struct Pkcs12 {
    pub identity: Identity,
    pub chain: Vec<Certificate>,
}

impl Pkcs12 {
    pub fn parse(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let dir = TempDir::new("native_tls").unwrap(); //fixme
        let keychain = try!(keychain::CreateOptions::new()
            .password(pass) // FIXME maybe generate a secure random password here?
            .create(dir.path().join("keychain")));

        let mut import = try!(Pkcs12ImportOptions::new()
            .passphrase(pass)
            .keychain(keychain)
            .import(buf));
        let import = import.pop().unwrap();

        // The identity's cert shows up in the chain, so filter it out to avoid sending twice
        // FIXME should probably use CFEquals here
        let identity_cert = try!(import.identity.certificate()).to_der();

        Ok(Pkcs12 {
            identity: Identity(import.identity),
            // FIXME possibly use the chain from the trust result instead?
            chain: import.cert_chain
                .into_iter()
                .filter(|c| c.to_der() != identity_cert)
                .map(Certificate)
                .collect(),
        })
    }
}

pub enum HandshakeError<S> {
    Interrupted(MidHandshakeTlsStream<S>),
    Failure(Error),
}

impl<S> From<secure_transport::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<base::Error> for HandshakeError<S> {
    fn from(e: base::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub struct MidHandshakeTlsStream<S>(secure_transport::MidHandshakeSslStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
    where S: fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S>
    where S: io::Read + io::Write
{
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct ClientBuilder(());

impl ClientBuilder {
    pub fn new() -> Result<ClientBuilder, Error> {
        Ok(ClientBuilder(()))
    }

    pub fn handshake<S>(&mut self,
                        domain: &str,
                        stream: S)
                        -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Client, ConnectionType::Stream));
        try!(ctx.set_peer_domain_name(domain));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct ServerBuilder {
    identity: SecIdentity,
    chain: Vec<SecCertificate>,
}

impl ServerBuilder {
    pub fn new<I>(identity: Identity, chain: I) -> Result<ServerBuilder, Error>
        where I: IntoIterator<Item = Certificate>
    {
        Ok(ServerBuilder {
            identity: identity.0,
            chain: chain.into_iter().map(|c| c.0).collect(),
        })
    }

    pub fn handshake<S>(&mut self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Server, ConnectionType::Stream));
        try!(ctx.set_certificate(&self.identity, &self.chain));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(secure_transport::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(try!(self.0.context().buffered_read_size()))
    }
}

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

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

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

/// Security Framework-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the Security Framework `SslStream`.
    fn raw_stream(&self) -> &secure_transport::SslStream<S>;

    /// Returns a mutable reference to the Security Framework `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &secure_transport::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S> {
        &mut (self.0).0
    }
}
<|endoftext|>


<commit_before>extern crate security_framework;
extern crate security_framework_sys;
extern crate tempdir;

use self::security_framework::base;
use self::security_framework::certificate::SecCertificate;
use self::security_framework::identity::SecIdentity;
use self::security_framework::import_export::Pkcs12ImportOptions;
use self::security_framework::secure_transport::{self, SslContext, ProtocolSide, ConnectionType};
use self::security_framework::secure_transport::SslProtocol;
use self::security_framework::os::macos::keychain;
use self::tempdir::TempDir;
use std::fmt;
use std::io;
use std::error;

pub struct Error(base::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<base::Error> for Error {
    fn from(error: base::Error) -> Error {
        Error(error)
    }
}

pub struct Pkcs12 {
    identity: SecIdentity,
    chain: Vec<SecCertificate>,
}

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let dir = match TempDir::new("native_tls") {
            Ok(dir) => dir,
            // Gotta throw away the real error :(
            Err(_) => return Err(Error(base::Error::from(security_framework_sys::base::errSecIO))),
        };

        let keychain = try!(keychain::CreateOptions::new()
            .password(pass) // FIXME maybe generate a secure random password here?
            .create(dir.path().join("keychain")));

        let mut import = try!(Pkcs12ImportOptions::new()
            .passphrase(pass)
            .keychain(keychain)
            .import(buf));
        let import = import.pop().unwrap();

        // The identity's cert shows up in the chain, so filter it out to avoid sending twice
        // FIXME should probably use CFEquals here
        let identity_cert = try!(import.identity.certificate()).to_der();

        Ok(Pkcs12 {
            identity: import.identity,
            // FIXME possibly use the chain from the trust result instead?
            chain: import.cert_chain
                .into_iter()
                .filter(|c| c.to_der() != identity_cert)
                .collect(),
        })
    }
}

pub enum HandshakeError<S> {
    Interrupted(MidHandshakeTlsStream<S>),
    Failure(Error),
}

impl<S> From<secure_transport::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<base::Error> for HandshakeError<S> {
    fn from(e: base::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub struct MidHandshakeTlsStream<S>(secure_transport::MidHandshakeSslStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
    where S: fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S>
    where S: io::Read + io::Write
{
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsConnectorBuilder(TlsConnector);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        self.0.pkcs12 = Some(pkcs12);
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(self.0)
    }
}

pub struct TlsConnector {
    pkcs12: Option<Pkcs12>,
}

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        Ok(TlsConnectorBuilder(TlsConnector {
            pkcs12: None,
        }))
    }

    pub fn connect<S>(&self,
                      domain: &str,
                      stream: S)
                      -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Client, ConnectionType::Stream));
        try!(ctx.set_peer_domain_name(domain));

        // Disable all protocols and then selectively re-enable as we can. This
        // is similar to what curl does on pre-10.8 builds.
        try!(ctx.set_protocol_version_enabled(SslProtocol::All, false));
        try!(ctx.set_protocol_version_enabled(SslProtocol::Tls1, true));
        if let Some(pkcs12) = self.pkcs12.as_ref() {
            try!(ctx.set_certificate(&pkcs12.identity, &pkcs12.chain));
        }
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsAcceptorBuilder(TlsAcceptor);

impl TlsAcceptorBuilder {
    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(self.0)
    }
}

pub struct TlsAcceptor {
    pkcs12: Pkcs12,
}

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        Ok(TlsAcceptorBuilder(TlsAcceptor {
            pkcs12: pkcs12,
        }))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Server, ConnectionType::Stream));
        try!(ctx.set_certificate(&self.pkcs12.identity, &self.pkcs12.chain));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(secure_transport::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(try!(self.0.context().buffered_read_size()))
    }
}

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

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

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

/// Security Framework-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the Security Framework `SslStream`.
    fn raw_stream(&self) -> &secure_transport::SslStream<S>;

    /// Returns a mutable reference to the Security Framework `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &secure_transport::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S> {
        &mut (self.0).0
    }
}
<commit_msg>Little cleanup<commit_after>extern crate security_framework;
extern crate security_framework_sys;
extern crate tempdir;

use self::security_framework::base;
use self::security_framework::certificate::SecCertificate;
use self::security_framework::identity::SecIdentity;
use self::security_framework::import_export::Pkcs12ImportOptions;
use self::security_framework::secure_transport::{self, SslContext, ProtocolSide, ConnectionType};
use self::security_framework::secure_transport::SslProtocol;
use self::security_framework::os::macos::keychain;
use self::security_framework_sys::base::errSecIO;
use self::tempdir::TempDir;
use std::fmt;
use std::io;
use std::error;

pub struct Error(base::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<base::Error> for Error {
    fn from(error: base::Error) -> Error {
        Error(error)
    }
}

pub struct Pkcs12 {
    identity: SecIdentity,
    chain: Vec<SecCertificate>,
}

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let dir = match TempDir::new("native_tls") {
            Ok(dir) => dir,
            // Gotta throw away the real error :(
            Err(_) => return Err(Error(base::Error::from(errSecIO))),
        };

        let keychain = try!(keychain::CreateOptions::new()
            .password(pass) // FIXME maybe generate a secure random password here?
            .create(dir.path().join("keychain")));

        let mut import = try!(Pkcs12ImportOptions::new()
            .passphrase(pass)
            .keychain(keychain)
            .import(buf));
        let import = import.pop().unwrap();

        // The identity's cert shows up in the chain, so filter it out to avoid sending twice
        // FIXME should probably use CFEquals here
        let identity_cert = try!(import.identity.certificate()).to_der();

        Ok(Pkcs12 {
            identity: import.identity,
            // FIXME possibly use the chain from the trust result instead?
            chain: import.cert_chain
                .into_iter()
                .filter(|c| c.to_der() != identity_cert)
                .collect(),
        })
    }
}

pub enum HandshakeError<S> {
    Interrupted(MidHandshakeTlsStream<S>),
    Failure(Error),
}

impl<S> From<secure_transport::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<base::Error> for HandshakeError<S> {
    fn from(e: base::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub struct MidHandshakeTlsStream<S>(secure_transport::MidHandshakeSslStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
    where S: fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S>
    where S: io::Read + io::Write
{
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsConnectorBuilder(TlsConnector);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        self.0.pkcs12 = Some(pkcs12);
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(self.0)
    }
}

pub struct TlsConnector {
    pkcs12: Option<Pkcs12>,
}

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        Ok(TlsConnectorBuilder(TlsConnector {
            pkcs12: None,
        }))
    }

    pub fn connect<S>(&self,
                      domain: &str,
                      stream: S)
                      -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Client, ConnectionType::Stream));
        try!(ctx.set_peer_domain_name(domain));

        // Disable all protocols and then selectively re-enable as we can. This
        // is similar to what curl does on pre-10.8 builds.
        try!(ctx.set_protocol_version_enabled(SslProtocol::All, false));
        try!(ctx.set_protocol_version_enabled(SslProtocol::Tls1, true));
        if let Some(pkcs12) = self.pkcs12.as_ref() {
            try!(ctx.set_certificate(&pkcs12.identity, &pkcs12.chain));
        }
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsAcceptorBuilder(TlsAcceptor);

impl TlsAcceptorBuilder {
    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(self.0)
    }
}

pub struct TlsAcceptor {
    pkcs12: Pkcs12,
}

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        Ok(TlsAcceptorBuilder(TlsAcceptor {
            pkcs12: pkcs12,
        }))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Server, ConnectionType::Stream));
        try!(ctx.set_certificate(&self.pkcs12.identity, &self.pkcs12.chain));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(secure_transport::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(try!(self.0.context().buffered_read_size()))
    }
}

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

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

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

/// Security Framework-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the Security Framework `SslStream`.
    fn raw_stream(&self) -> &secure_transport::SslStream<S>;

    /// Returns a mutable reference to the Security Framework `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &secure_transport::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S> {
        &mut (self.0).0
    }
}
<|endoftext|>


<commit_before>extern crate security_framework;
extern crate security_framework_sys;

use self::security_framework::base;
use self::security_framework::certificate::SecCertificate;
use self::security_framework::identity::SecIdentity;
use self::security_framework::import_export::Pkcs12ImportOptions;
use self::security_framework::secure_transport::{self, SslContext, ProtocolSide, ConnectionType,
                                                 SslProtocol, ClientBuilder};
use std::fmt;
use std::io;
use std::error;

use Protocol;

fn convert_protocol(protocol: Protocol) -> SslProtocol {
    match protocol {
        Protocol::Sslv3 => SslProtocol::Ssl3,
        Protocol::Tlsv10 => SslProtocol::Tls1,
        Protocol::Tlsv11 => SslProtocol::Tls11,
        Protocol::Tlsv12 => SslProtocol::Tls12,
        Protocol::__NonExhaustive => unreachable!(),
    }
}

fn protocol_min_max(protocols: &[Protocol]) -> (SslProtocol, SslProtocol) {
    let mut min = Protocol::Tlsv12;
    let mut max = Protocol::Sslv3;
    for protocol in protocols {
        if (*protocol as usize) < (min as usize) {
            min = *protocol;
        }
        if (*protocol as usize) > (max as usize) {
            max = *protocol;
        }
    }
    (convert_protocol(min), convert_protocol(max))
}

pub struct Error(base::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<base::Error> for Error {
    fn from(error: base::Error) -> Error {
        Error(error)
    }
}

pub struct Pkcs12 {
    identity: SecIdentity,
    chain: Vec<SecCertificate>,
}

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let mut options = Pkcs12ImportOptions::new();
        options.passphrase(pass);

        let mut imported_identities = try!(options.import(buf));
        imported_identities.truncate(1);
        let imported_identity = imported_identities.pop().unwrap();

        // FIXME: Compare the certificates for equality using CFEqual
        let identity_cert = try!(imported_identity.identity.certificate()).to_der();

        Ok(Pkcs12{
            identity: imported_identity.identity,
            chain: imported_identity.cert_chain
                .into_iter()
                .filter(|c| c.to_der() != identity_cert)
                .collect(),})
    }
}

pub enum HandshakeError<S> {
    Interrupted(MidHandshakeTlsStream<S>),
    Failure(Error),
}

impl<S> From<secure_transport::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream::Server(s))
            }
        }
    }
}

impl<S> From<secure_transport::ClientHandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::ClientHandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::ClientHandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::ClientHandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream::Client(s))
            }
        }
    }
}

impl<S> From<base::Error> for HandshakeError<S> {
    fn from(e: base::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub enum MidHandshakeTlsStream<S> {
    Server(secure_transport::MidHandshakeSslStream<S>),
    Client(secure_transport::MidHandshakeClientBuilder<S>),
}

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
    where S: fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MidHandshakeTlsStream::Server(ref s) => s.fmt(fmt),
            MidHandshakeTlsStream::Client(ref s) => s.fmt(fmt),
        }
    }
}

impl<S> MidHandshakeTlsStream<S>
    where S: io::Read + io::Write
{
    pub fn get_ref(&self) -> &S {
        match *self {
            MidHandshakeTlsStream::Server(ref s) => s.get_ref(),
            MidHandshakeTlsStream::Client(ref s) => s.get_ref(),
        }
    }

    pub fn get_mut(&mut self) -> &mut S {
        match *self {
            MidHandshakeTlsStream::Server(ref mut s) => s.get_mut(),
            MidHandshakeTlsStream::Client(ref mut s) => s.get_mut(),
        }
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self {
            MidHandshakeTlsStream::Server(s) => {
                match s.handshake() {
                    Ok(s) => Ok(TlsStream(s)),
                    Err(e) => Err(e.into()),
                }
            }
            MidHandshakeTlsStream::Client(s) => {
                match s.handshake() {
                    Ok(s) => Ok(TlsStream(s)),
                    Err(e) => Err(e.into()),
                }
            }
        }
    }
}

pub struct TlsConnectorBuilder(TlsConnector);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        self.0.pkcs12 = Some(pkcs12);
        Ok(())
    }

    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        self.0.protocols = protocols.to_vec();
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(self.0)
    }
}

pub struct TlsConnector {
    pkcs12: Option<Pkcs12>,
    protocols: Vec<Protocol>,
    anchor_certificates: Option<Vec<SecCertificate>>,
}

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        Ok(TlsConnectorBuilder(TlsConnector {
            pkcs12: None,
            protocols: vec![Protocol::Tlsv10, Protocol::Tlsv11, Protocol::Tlsv12],
            anchor_certificates: None,
        }))
    }

    pub fn connect<S>(&self,
                      domain: &str,
                      stream: S)
                      -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut builder = ClientBuilder::new();
        let (min, max) = protocol_min_max(&self.protocols);
        builder.protocol_min(min);
        builder.protocol_max(max);
        if let Some(pkcs12) = self.pkcs12.as_ref() {
            builder.identity(&pkcs12.identity, &pkcs12.chain);
        }
        if let Some(anchors) = self.anchor_certificates.as_ref() {
            builder.anchor_certificates(anchors);
        }
        match builder.handshake2(domain, stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsAcceptorBuilder(TlsAcceptor);

impl TlsAcceptorBuilder {
    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        self.0.protocols = protocols.to_vec();
        Ok(())
    }

    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(self.0)
    }
}

pub struct TlsAcceptor {
    pkcs12: Pkcs12,
    protocols: Vec<Protocol>,
}

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        Ok(TlsAcceptorBuilder(TlsAcceptor {
            pkcs12: pkcs12,
            protocols: vec![Protocol::Tlsv10, Protocol::Tlsv11, Protocol::Tlsv12],
        }))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Server, ConnectionType::Stream));

        let (min, max) = protocol_min_max(&self.protocols);
        try!(ctx.set_protocol_version_min(min));
        try!(ctx.set_protocol_version_max(max));
        try!(ctx.set_certificate(&self.pkcs12.identity, &self.pkcs12.chain));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(secure_transport::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(try!(self.0.context().buffered_read_size()))
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        try!(self.0.close());
        Ok(())
    }
}

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

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

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

/// Security Framework-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the Security Framework `SslStream`.
    fn raw_stream(&self) -> &secure_transport::SslStream<S>;

    /// Returns a mutable reference to the Security Framework `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &secure_transport::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S> {
        &mut (self.0).0
    }
}

/// Security Framework-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Specifies the set of additional root certificates to trust when
    /// verifying the server's certificate.
    fn anchor_certificates(&mut self, certs: &[SecCertificate]) -> &mut Self;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn anchor_certificates(&mut self, certs: &[SecCertificate]) -> &mut Self {
        (self.0).0.anchor_certificates = Some(certs.to_owned());
        self
    }
}

/// Security Framework-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying Security Framework error for inspection.
    fn security_framework_error(&self) -> &base::Error;
}

impl ErrorExt for ::Error {
    fn security_framework_error(&self) -> &base::Error {
        &(self.0).0
    }
}
<commit_msg>Little formatting cleanup<commit_after>extern crate security_framework;
extern crate security_framework_sys;

use self::security_framework::base;
use self::security_framework::certificate::SecCertificate;
use self::security_framework::identity::SecIdentity;
use self::security_framework::import_export::Pkcs12ImportOptions;
use self::security_framework::secure_transport::{self, SslContext, ProtocolSide, ConnectionType,
                                                 SslProtocol, ClientBuilder};
use std::fmt;
use std::io;
use std::error;

use Protocol;

fn convert_protocol(protocol: Protocol) -> SslProtocol {
    match protocol {
        Protocol::Sslv3 => SslProtocol::Ssl3,
        Protocol::Tlsv10 => SslProtocol::Tls1,
        Protocol::Tlsv11 => SslProtocol::Tls11,
        Protocol::Tlsv12 => SslProtocol::Tls12,
        Protocol::__NonExhaustive => unreachable!(),
    }
}

fn protocol_min_max(protocols: &[Protocol]) -> (SslProtocol, SslProtocol) {
    let mut min = Protocol::Tlsv12;
    let mut max = Protocol::Sslv3;
    for protocol in protocols {
        if (*protocol as usize) < (min as usize) {
            min = *protocol;
        }
        if (*protocol as usize) > (max as usize) {
            max = *protocol;
        }
    }
    (convert_protocol(min), convert_protocol(max))
}

pub struct Error(base::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<base::Error> for Error {
    fn from(error: base::Error) -> Error {
        Error(error)
    }
}

pub struct Pkcs12 {
    identity: SecIdentity,
    chain: Vec<SecCertificate>,
}

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let mut options = Pkcs12ImportOptions::new();
        options.passphrase(pass);

        let mut imported_identities = try!(options.import(buf));
        imported_identities.truncate(1);
        let imported_identity = imported_identities.pop().unwrap();

        // FIXME: Compare the certificates for equality using CFEqual
        let identity_cert = try!(imported_identity.identity.certificate()).to_der();

        Ok(Pkcs12 {
            identity: imported_identity.identity,
            chain: imported_identity.cert_chain
                .into_iter()
                .filter(|c| c.to_der() != identity_cert)
                .collect(),
        })
    }
}

pub enum HandshakeError<S> {
    Interrupted(MidHandshakeTlsStream<S>),
    Failure(Error),
}

impl<S> From<secure_transport::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream::Server(s))
            }
        }
    }
}

impl<S> From<secure_transport::ClientHandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::ClientHandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::ClientHandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::ClientHandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream::Client(s))
            }
        }
    }
}

impl<S> From<base::Error> for HandshakeError<S> {
    fn from(e: base::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub enum MidHandshakeTlsStream<S> {
    Server(secure_transport::MidHandshakeSslStream<S>),
    Client(secure_transport::MidHandshakeClientBuilder<S>),
}

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
    where S: fmt::Debug
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MidHandshakeTlsStream::Server(ref s) => s.fmt(fmt),
            MidHandshakeTlsStream::Client(ref s) => s.fmt(fmt),
        }
    }
}

impl<S> MidHandshakeTlsStream<S>
    where S: io::Read + io::Write
{
    pub fn get_ref(&self) -> &S {
        match *self {
            MidHandshakeTlsStream::Server(ref s) => s.get_ref(),
            MidHandshakeTlsStream::Client(ref s) => s.get_ref(),
        }
    }

    pub fn get_mut(&mut self) -> &mut S {
        match *self {
            MidHandshakeTlsStream::Server(ref mut s) => s.get_mut(),
            MidHandshakeTlsStream::Client(ref mut s) => s.get_mut(),
        }
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self {
            MidHandshakeTlsStream::Server(s) => {
                match s.handshake() {
                    Ok(s) => Ok(TlsStream(s)),
                    Err(e) => Err(e.into()),
                }
            }
            MidHandshakeTlsStream::Client(s) => {
                match s.handshake() {
                    Ok(s) => Ok(TlsStream(s)),
                    Err(e) => Err(e.into()),
                }
            }
        }
    }
}

pub struct TlsConnectorBuilder(TlsConnector);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        self.0.pkcs12 = Some(pkcs12);
        Ok(())
    }

    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        self.0.protocols = protocols.to_vec();
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(self.0)
    }
}

pub struct TlsConnector {
    pkcs12: Option<Pkcs12>,
    protocols: Vec<Protocol>,
    anchor_certificates: Option<Vec<SecCertificate>>,
}

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        Ok(TlsConnectorBuilder(TlsConnector {
            pkcs12: None,
            protocols: vec![Protocol::Tlsv10, Protocol::Tlsv11, Protocol::Tlsv12],
            anchor_certificates: None,
        }))
    }

    pub fn connect<S>(&self,
                      domain: &str,
                      stream: S)
                      -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut builder = ClientBuilder::new();
        let (min, max) = protocol_min_max(&self.protocols);
        builder.protocol_min(min);
        builder.protocol_max(max);
        if let Some(pkcs12) = self.pkcs12.as_ref() {
            builder.identity(&pkcs12.identity, &pkcs12.chain);
        }
        if let Some(anchors) = self.anchor_certificates.as_ref() {
            builder.anchor_certificates(anchors);
        }
        match builder.handshake2(domain, stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsAcceptorBuilder(TlsAcceptor);

impl TlsAcceptorBuilder {
    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        self.0.protocols = protocols.to_vec();
        Ok(())
    }

    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(self.0)
    }
}

pub struct TlsAcceptor {
    pkcs12: Pkcs12,
    protocols: Vec<Protocol>,
}

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        Ok(TlsAcceptorBuilder(TlsAcceptor {
            pkcs12: pkcs12,
            protocols: vec![Protocol::Tlsv10, Protocol::Tlsv11, Protocol::Tlsv12],
        }))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ctx = try!(SslContext::new(ProtocolSide::Server, ConnectionType::Stream));

        let (min, max) = protocol_min_max(&self.protocols);
        try!(ctx.set_protocol_version_min(min));
        try!(ctx.set_protocol_version_max(max));
        try!(ctx.set_certificate(&self.pkcs12.identity, &self.pkcs12.chain));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(secure_transport::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(try!(self.0.context().buffered_read_size()))
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        try!(self.0.close());
        Ok(())
    }
}

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

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

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

/// Security Framework-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the Security Framework `SslStream`.
    fn raw_stream(&self) -> &secure_transport::SslStream<S>;

    /// Returns a mutable reference to the Security Framework `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &secure_transport::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S> {
        &mut (self.0).0
    }
}

/// Security Framework-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Specifies the set of additional root certificates to trust when
    /// verifying the server's certificate.
    fn anchor_certificates(&mut self, certs: &[SecCertificate]) -> &mut Self;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn anchor_certificates(&mut self, certs: &[SecCertificate]) -> &mut Self {
        (self.0).0.anchor_certificates = Some(certs.to_owned());
        self
    }
}

/// Security Framework-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying Security Framework error for inspection.
    fn security_framework_error(&self) -> &base::Error;
}

impl ErrorExt for ::Error {
    fn security_framework_error(&self) -> &base::Error {
        &(self.0).0
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! How and where to record the Serde format of interesting Libra types.
//! See API documentation with `cargo doc -p serde-reflection --open`

use serde_reflection::Registry;
use std::str::FromStr;
use structopt::{clap::arg_enum, StructOpt};

/// Consensus messages.
mod consensus;
/// Libra transactions.
mod libra;
/// Analyze Serde formats to detect certain patterns.
mod linter;
/// Move ABI.
mod move_abi;
/// Network messages.
mod network;

pub use linter::lint_lcs_format;

arg_enum! {
#[derive(Debug, StructOpt, Clone, Copy)]
/// A corpus of Rust types to trace, and optionally record on disk.
pub enum Corpus {
    Libra,
    Consensus,
    Network,
    MoveABI,
}
}

impl Corpus {
    /// All corpuses.
    pub fn values() -> Vec<Corpus> {
        Corpus::variants()
            .iter()
            .filter_map(|s| Corpus::from_str(s).ok())
            .collect()
    }

    /// Compute the registry of formats.
    pub fn get_registry(self) -> Registry {
        match self {
            Corpus::Libra => libra::get_registry().unwrap(),
            Corpus::Consensus => consensus::get_registry().unwrap(),
            Corpus::Network => network::get_registry().unwrap(),
            Corpus::MoveABI => move_abi::get_registry().unwrap(),
        }
    }

    /// Where to record this corpus on disk.
    pub fn output_file(self) -> Option<&'static str> {
        match self {
            Corpus::Libra => libra::output_file(),
            Corpus::Consensus => consensus::output_file(),
            Corpus::Network => network::output_file(),
            Corpus::MoveABI => move_abi::output_file(),
        }
    }
}
<commit_msg>[generate format] upgrade to serde-reflection 0.3.1 for better error reporting<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! How and where to record the Serde format of interesting Libra types.
//! See API documentation with `cargo doc -p serde-reflection --open`

use serde_reflection::Registry;
use std::str::FromStr;
use structopt::{clap::arg_enum, StructOpt};

/// Consensus messages.
mod consensus;
/// Libra transactions.
mod libra;
/// Analyze Serde formats to detect certain patterns.
mod linter;
/// Move ABI.
mod move_abi;
/// Network messages.
mod network;

pub use linter::lint_lcs_format;

arg_enum! {
#[derive(Debug, StructOpt, Clone, Copy)]
/// A corpus of Rust types to trace, and optionally record on disk.
pub enum Corpus {
    Libra,
    Consensus,
    Network,
    MoveABI,
}
}

impl Corpus {
    /// All corpuses.
    pub fn values() -> Vec<Corpus> {
        Corpus::variants()
            .iter()
            .filter_map(|s| Corpus::from_str(s).ok())
            .collect()
    }

    /// Compute the registry of formats.
    pub fn get_registry(self) -> Registry {
        let result = match self {
            Corpus::Libra => libra::get_registry(),
            Corpus::Consensus => consensus::get_registry(),
            Corpus::Network => network::get_registry(),
            Corpus::MoveABI => move_abi::get_registry(),
        };
        match result {
            Ok(registry) => registry,
            Err(error) => {
                panic!("{}:{}", error, error.explanation());
            }
        }
    }

    /// Where to record this corpus on disk.
    pub fn output_file(self) -> Option<&'static str> {
        match self {
            Corpus::Libra => libra::output_file(),
            Corpus::Consensus => consensus::output_file(),
            Corpus::Network => network::output_file(),
            Corpus::MoveABI => move_abi::output_file(),
        }
    }
}
<|endoftext|>